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

2026-08-22 赤色 Master / Phase 8+ ★★★★★★★★★ forced/excludedパーティション+優先度付きキュー

問題

$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

3 3 2
0 1 1
1 2 2
0 2 3

出力例1

4

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

概念図

パーティション分割: forced/excludedで重複なく子を生成 F=∅,E=∅ → MST w=3 F=∅,E={e1=01} → w=4 F={e1},E={e2=12} → w=5 ヒープから軽い順に取り出し、その使用辺(F以外)を1本ずつ除外する子パーティションを生成 子パーティション同士は互いに素で、元の全域木以外を1回ずつ過不足なくカバーする

ヒント(段階的開示)

ヒント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))

自己評価

自分の回答

気づき・メモ