Day 130-Q4 | K番目最小全域木（Lawler法によるパーティション分割探索）

問題

$N$頂点$M$辺の無向連結グラフが与えられる。全域木（異なる辺集合として区別し重みが同じでも別々に数える）を総重みの小さい順に並べたとき、$K$番目に小さい総重みを求めよ。

入力形式

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

制約

$2 \le N \le 500$

$N-1 \le M \le 2000$

$1 \le K \le 100$

グラフは連結、K番目の全域木の存在を保証

入出力例

入力例1

出力例1

全域木は{01,12}重み3(MST)、{01,02}重み4、{12,02}重み5の3通り。2番目は4。

概念図

ヒント（段階的開示）

ヒント1（方向性）

2番目に軽い全域木はMSTの辺を1本交換して得られる。一般のKに拡張するには、候補全域木の集合をもれなく分割し尽くす枠組みが必要。

ヒント2（アプローチ）

状態を「強制的に使う辺集合F」と「絶対使わない辺集合E」の組(パーティション)で表す。パーティションの最小全域木はFを先に使い、残りをKruskal法でEを除いて貪欲追加すれば求まる。ヒープから軽い順に取り出す度に使用辺を1本ずつ除外する子パーティションへ重複なく分割する(Lawler法/Murtyの一般化)。

ヒント3（誘導）

tree_edges = [e for e in used if e not in forced]
for i, e in enumerate(tree_edges):
    new_forced = forced | set(tree_edges[:i])
    new_excluded = excluded | {e}
    res = find_mst(n, edges, new_forced, new_excluded)
    if res is not None:
        heappush(heap, (res[0], new_forced, new_excluded, res[1]))

模範解答 (Python)

import sys
import heapq


class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        self.parent[ra] = rb
        return True


def find_mst(n, edges, order, forced, excluded):
    uf = UnionFind(n)
    total = 0
    used = []
    for idx in forced:
        u, v, w = edges[idx]
        if not uf.union(u, v):
            return None
        used.append(idx)
        total += w
    for idx in order:
        if idx in forced or idx in excluded:
            continue
        u, v, w = edges[idx]
        if uf.union(u, v):
            used.append(idx)
            total += w
    if len(used) != n - 1:
        return None
    return total, used


def kth_mst(n, edges, k):
    order = sorted(range(len(edges)), key=lambda i: edges[i][2])
    base = find_mst(n, edges, order, frozenset(), frozenset())
    heap = []
    counter = 0
    heapq.heappush(heap, (base[0], counter, frozenset(), frozenset(), tuple(base[1])))
    counter += 1
    found = 0
    while heap:
        w, _, forced, excluded, used = heapq.heappop(heap)
        found += 1
        if found == k:
            return w
        tree_edges = [e for e in used if e not in forced]
        for i, e in enumerate(tree_edges):
            new_forced = forced | frozenset(tree_edges[:i])
            new_excluded = excluded | frozenset([e])
            res = find_mst(n, edges, order, new_forced, new_excluded)
            if res is not None:
                heapq.heappush(heap, (res[0], counter, new_forced, new_excluded, tuple(res[1])))
                counter += 1
    return -1


def main():
    data = sys.stdin.read().split()
    idx = 0
    n = int(data[idx]); m = int(data[idx + 1]); k = int(data[idx + 2]); idx += 3
    edges = []
    for _ in range(m):
        u = int(data[idx]); v = int(data[idx + 1]); w = int(data[idx + 2]); idx += 3
        edges.append((u, v, w))

    print(kth_mst(n, edges, k))


main()

Step-by-Step 解説

1パーティションという発想
全ての全域木を「強制辺F・除外辺E」ごとに互いに重複しないグループへ分割し、各グループの最小全域木を代表として比較する。

2最初のパーティション
F=∅,E=∅から始める。この最小全域木は通常のKruskal法によるMSTそのもの。

3重複なく分割する工夫
使用辺e1..emを順に「e1除外」「e1使用,e2除外」...と分割すると、子パーティションは互いに素で元の代表以外の全域木を過不足なくカバーする。

4ヒープでの取り出し
軽い順に取り出すたびK番目候補が確定し、同時に子パーティションを生成してヒープに積む。K回目の取り出しが答え。

5計算量
1回の分割で高々O(N)個の子を生成し各子でO(Mα(N))のKruskal法。全体でO(K N M α(N))程度。

よくあるミス

ミス	原因	正しい書き方
子パーティション生成の順序基準を統一しない	分割の非重複性は同じ順序基準に依存する	辺重みでソートしたorderに基づき常に同じ基準で分割する
forced集合が閉路を作る判定を忘れる	unionが失敗した時点で終了しないと不正な木を数える	unionがFalseを返したら直ちにNoneを返し実行不可能扱いにする
ヒープの比較でset同士を比較して例外	同じ優先度の要素同士の比較が発生	タイブレーク用の一意カウンタを2番目要素に入れる
Kが全域木総数を超える入力を想定しない	制約を読み飛ばす	「K番目が存在することを保証」を前提にする

次のステップ

発展: 同じ発想はK番目最短路(Yen's algorithm)やK番目最小重み完全マッチングなど、他のマトロイド的最適化問題にも応用できる
次回予告: 最近点対問題（分割統治法によるO(N log N)）

自己評価

完全理解おおむね理解要復習

Day 130-Q4 — K番目最小全域木（Lawler法によるパーティション分割探索）

問題