Day 124-Q3 — Weighted Matroid Union（重み付きマトロイド合併・K-forest分割）

問題

$N$ 頂点 $M$ 本の重み付き無向辺を持つグラフが与えられる。整数 $K$ が指定される。

辺の部分集合 $S$ を選び、$S$ を高々 $K$ 本の互いに辺素な森（forest）に分割できる（＝どの1本の森も閉路を含まない）とき、$S$ は「$K$-森分割可能」であるという。

$K$-森分割可能な部分集合 $S$ の中で、辺の重みの総和が最大になるものを1つ求め、その最大重み総和と、各採用辺がどの森（$1$〜$K$）に属するかを出力せよ。

入力形式

N M K
u_1 v_1 w_1
...
u_M v_M w_M

制約

入出力例

4 5 1
1 2 5
2 3 4
3 4 3
4 1 2
1 3 6

14

概念図

ヒント（段階的開示）

def can_place_directly(forest_id, u, v, colors):
    # forest_id の辺集合だけでDSUを作りu,vが非連結かチェック
    ...

def find_augmenting(e, colors, K):
    # colors[edge] = 現在の森番号 or None
    # BFSで e を起点に「eをforest iに入れたときfをどかせば良い」の連鎖を探す
    ...

模範解答 (Python)

📋 コピーimport sys
from itertools import combinations

def solve():
    data = sys.stdin.read().split()
    idx = 0
    N = int(data[idx]); idx += 1
    M = int(data[idx]); idx += 1
    K = int(data[idx]); idx += 1
    edges = []
    for i in range(M):
        u = int(data[idx]); idx += 1
        v = int(data[idx]); idx += 1
        w = int(data[idx]); idx += 1
        edges.append((u, v, w))

    class DSU:
        def __init__(self, n):
            self.p = list(range(n + 1))
        def find(self, x):
            while self.p[x] != x:
                self.p[x] = self.p[self.p[x]]
                x = self.p[x]
            return x
        def union(self, x, y):
            rx, ry = self.find(x), self.find(y)
            if rx == ry:
                return False
            self.p[rx] = ry
            return True

    def independent_in(forest_edges_idx, excluded, extra):
        # forest_edges_idx: 現在その森に属している辺インデックス集合
        # excluded: 除外する辺インデックス（付け替えで一時的に抜く）
        # extra: 追加で入れてみる (u, v)
        dsu = DSU(N)
        for ei in forest_edges_idx:
            if ei == excluded:
                continue
            u, v, w = edges[ei]
            if u == v:
                return False
            dsu.union(u, v)
        u, v = extra
        if u == v:
            return False
        return dsu.union(u, v)

    color = [None] * M  # color[i] = 0..K-1 or None
    forest_members = [[] for _ in range(K)]

    order = sorted(range(M), key=lambda i: -edges[i][2])

    for e in order:
        u, v, w = edges[e]
        if u == v:
            continue  # 自己ループはどの森にも入らない
        # 直接入れられる森を探す
        placed = False
        for c in range(K):
            if independent_in(forest_members[c], None, (u, v)):
                color[e] = c
                forest_members[c].append(e)
                placed = True
                break
        if placed:
            continue

        # 交換増加路探索: BFS on edges
        # state: 現在「まだ居場所が決まっていない辺」= cur (最初は e)
        # 各 cur について、各森 c 内の辺 f を1つ抜いて cur を入れられるなら (c, f) が候補
        # f はその後「新たに居場所を探す辺」になる
        visited_edges = {e}
        parent = {}   # child_edge -> (forest_c, replaced_edge_f) の記録
        queue = [e]
        found = None
        while queue and found is None:
            cur = queue.pop(0)
            cu, cv, cw = edges[cur]
            for c in range(K):
                # cur を forest c に「直接」入れられるか（別の辺を抜かずに）
                if independent_in(forest_members[c], None, (cu, cv)):
                    found = (cur, c, None)
                    break
                # forest c 内の各辺 f を抜いて cur を入れられるか調べる
                for f in forest_members[c]:
                    if independent_in(forest_members[c], f, (cu, cv)):
                        if f not in visited_edges:
                            visited_edges.add(f)
                            parent[f] = (cur, c)
                            queue.append(f)
            if found:
                break

        if found is None:
            continue  # これ以上 e を追加できない

        # 経路復元して付け替えを実行
        cur, final_color, _ = found
        # cur から e まで parent を遡って、玉突きの連鎖を構築
        chain = [cur]
        node = cur
        while node != e:
            prev, prev_color = parent[node]
            chain.append(prev)
            node = prev
        chain.reverse()  # chain[0] == e, chain[-1] == cur (居場所が確定する辺)

        # 実際の付け替え: chain[i] は元々 chain[i+1] がいた森から追い出されて、
        # chain[i] 自身は次段階で別の森 (parent[chain[i+1]] の c) に入る、最終的に chain[-1] は final_color に直接入る
        # 簡潔にするため、各辺の新しい色を後ろから確定させる
        new_color_for = {}
        new_color_for[chain[-1]] = final_color
        for i in range(len(chain) - 2, -1, -1):
            nxt = chain[i + 1]
            c_of_nxt_origin = parent[nxt][1]  # nxt が元々いた森
            new_color_for[chain[i]] = c_of_nxt_origin

        for ei in chain:
            old_c = color[ei]
            if old_c is not None:
                forest_members[old_c].remove(ei)
            new_c = new_color_for[ei]
            color[ei] = new_c
            forest_members[new_c].append(ei)

    total = sum(edges[i][2] for i in range(M) if color[i] is not None)
    print(total)
    result_lines = []
    for i in range(M):
        if color[i] is not None:
            result_lines.append(f"{i} {color[i] + 1}")
    for line in result_lines:
        print(line)

solve()

Step-by-Step 解説

よくあるミス

次のステップ

• 発展: 一般の（グラフィックとは限らない）マトロイド同士の合併では、独立性オラクルが与えられている前提でこの交換増加路アルゴリズムがそのまま使える。分割マトロイド・線形マトロイドとの合併に一般化すると、スケジューリング問題や資源割当問題に応用できる。
• 次回予告: Eppstein's K-Shortest Paths Algorithm（サイドトラック辺とpersistent leftist heapによるk最短ウォーク列挙）

自己評価

自分の回答

気づき・メモ