[BOJ 2887, Python] 惑星トンネル

問題リンク

BOJ 2887

分類

グラフ理論、ソート、最小全域木

解説

この問題に最初に出会ってから、もう半年以上経っていることに驚かされた…とはいえ最近プラチナ問題を疎かにしていた割には、それでも納得できる解法で解けたという点で意味は十分だったと思い、記録することにした。

この問題を正攻法で解くには、最大 10 万個の座標から 2 つを選び、その値を最小値から並べた上で最小辺から取り込んでいく Union-Find + Kruskal の組み合わせで解ける。

しかしそうやって解くと 10⁵ * 10⁵ なので、軽く 100 億通りの場合の数になってしまい、結果として 100 秒程度の実行時間が必要になるため、不可能な解法と言える。

Tip: 1 億回の演算につき約 1 秒と見積もると、計算量を考えやすい！

したがってこの問題の鍵は、O(N²) の場合の数をどう O(N) 単位まで下げるか、と言える。

結論から言うと、本問題では任意の点について全域木を構築するとき、隣接する点以外の点は考慮しなくてよい。

任意の点 A, B が存在し、C という点に接続しなければならない状況を見る。

x の差分 (|X_A - X_B|) が y, z より小さい最小値である場合 → C という点の x 値が A と B の x 値の間に存在すると仮定する

1. |X_A - X_C| + |X_B - X_C| となる場合、差は存在しない
2. |X_A - X_B| を維持したまま C を追加でつなぐ任意の |X_A - X_C| や |X_B - X_C| が存在するなら必ず既存の全域木より値が大きくなる → 反例の存在により不可

したがって任意の 2 点を結ぶ辺が本問題で最小全域木に含まれるなら、その 2 点は必ず隣接していなければならない、という証明になる！

この証明を理解できれば鍵を押さえたことになり、残りは一般的な Union-Find 問題と大きく変わらないと思う。

---

追記) heapq の代わりに PriorityQueue を使ったとき時間切れが発生したが、その問題への回答を見つけることができた。

What's the difference between heapq and PriorityQueue in python?

PriorityQueue は Thread Safety を保証するため遅いとのこと。

解法コード

# 惑星トンネル

import sys
from heapq import heappush, heappop

input = sys.stdin.readline

def check_union(union_info: list, now_target: int) -> int:
    """
    該当ノードの現在の Union を照会する関数
    """

    if union_info[now_target] == -1:
        union_info[now_target] = now_target

    elif union_info[now_target] != now_target:
        union_info[now_target] = check_union(union_info, union_info[now_target])

    return union_info[now_target]

def make_union(union_info: list, target_1: int, target_2: int) -> None:
    """
    2 点が同じ Union か確認し、違う場合に統合する関数
    """
    union_1 = check_union(union_info, target_1)
    union_2 = check_union(union_info, target_2)

    if union_1 == union_2:
        return True

    else:
        if union_1 > union_2:
            union_info[union_1] = union_2
        else:
            union_info[union_2] = union_1

        return False

planet_number = int(input())
coord_list = []

for i in range(planet_number):
    coord_list.append([i] + list(map(int, input().split())))

cost_queue = []

for i in range(1, 4):
    coord_list.sort(key=lambda x: x[i])
    for j in range(planet_number - 1):
        heappush(
            cost_queue,
            (
                coord_list[j + 1][i] - coord_list[j][i],
                (coord_list[j][0], coord_list[j + 1][0]),
            ),
        )

union_info = [-1 for _ in range(planet_number)]

total_length = 0
line_number = 0

while cost_queue:
    if line_number == planet_number - 1:
        break

    now_length, nodes = heappop(cost_queue)
    node_a, node_b = nodes

    if not make_union(union_info, node_a, node_b):
        total_length += now_length
        line_number += 1

print(total_length)