Pythonで森（フォレスト）を1本の木に接続するプログラム

隣接リスト形式でグラフが与えられているとします。このグラフは実際には、互いに連結していない複数の木から構成される「森（フォレスト）」です。ここで、いくつかの辺を追加して森全体を1本の木につなげることを考えます。その際、任意の2つのノード間の最長経路の距離（木の直径）が最小になるようにしなければなりません。

例えば、次のような入力が与えられた場合を考えてみます。

このとき、出力は 4 になります。

具体的には、ノード0とノード5の間に辺を追加すると、最長経路は「3 → 1 → 0 → 5 → 7」あるいは「4 → 1 → 0 → 5 → 7」（およびその逆方向の経路）のいずれかになります。したがって、答えは距離4です。

解法のアプローチ

この問題の鍵となるのは、各木の「直径」と「半径」の関係です。木の直径を d とすると、その木の中心から最も遠い葉までの距離は ⌈d / 2⌉ になります。木同士を中心同士でつなぐように辺を追加すれば、完成した木の直径を最小にできます。また、どのように辺を追加しても、元の森の中で最も直径の大きかった木より短くすることはできないため、その値も答えの候補として保持しておく必要があります。

これを踏まえて、以下の手順で問題を解きます。

• seen := 空のセット（訪問済みノードの管理用）、dic := グラフ（隣接リスト）

• 関数 treeDepth(node) を定義します。ノードを引数にとり、そのノードが属する木の直径を返します。

  • ret := 0

  • 関数 dfs1(node, parent) を定義します。

    • node を seen に追加する

    • best2 := 空の最小ヒープ

    • dic[node] 内の各隣接ノード nxt に対して:

      • nxt が parent と異なる場合、dfs1(nxt, node) + 1 を best2 にプッシュする

    • best2 のサイズが2を超えたら、最小要素をポップする（常に上位2つの深さだけを保持）

    • best2 が空の場合（葉ノードの場合）は 0 を返す

    • ret := max(ret, sum(best2))（このノードを頂点とする最長パスで直径を更新）

    • max(best2) を返す（親には最も深い1本だけを報告）

  • dfs1(node, None) を呼び出す

  • ret を返す

• メイン処理では以下を行います。

  • ret := 0、opt := 空のリスト、sing := 0

  • 0 からグラフのサイズ未満の各ノードについて:

    • node が seen に含まれていれば、次の反復へ進む

    • res := treeDepth(node)

    • sing := max(sing, res)（既存の木の中での最大直径を記録）

    • opt の末尾に ⌈res / 2⌉ を追加する

  • opt のサイズが1以下であれば、sing を返す（森がすでに1本の木である場合）

  • mx := max(opt)

  • opt 内で mx と等しい最初の要素を1減らす（最大の木の中心を接続基準に合わせる調整）

  • opt のすべての要素に1を加える（接続に使用する辺1本分を反映）

  • high2 := opt の大きい方から2つの要素

  • max(sum(high2), sing) を返す

実装例

それでは、実際のPythonコードを見てみましょう。

import heapq, math
class Solution:
    def solve(self, graph):
        seen = set()
        dic = graph
        def treeDepth(node):
            self.ret = 0
            def dfs1(node, parent):
                seen.add(node)
                best2 = []
                for nxt in dic[node]:
                    if nxt != parent:
                        heapq.heappush(best2, dfs1(nxt, node) + 1)
                        if len(best2) > 2:
                            heapq.heappop(best2)
                if not best2:
                    return 0
                self.ret = max(self.ret, sum(best2))
                return max(best2)
            dfs1(node, None)
            return self.ret
        ret = 0
        opt = []
        sing = 0
        for node in range(len(graph)):
            if node in seen:
                continue
            res = treeDepth(node)
            sing = max(sing, res)
            opt.append(int(math.ceil(res / 2)))
        if len(opt) <= 1:
            return sing
        mx = max(opt)
        for i in range(len(opt)):
            if opt[i] == mx:
                opt[i] -= 1
                break
        for i in range(len(opt)):
            opt[i] += 1
        high2 = heapq.nlargest(2, opt)
        return max(sum(high2), sing)
ob = Solution()
graph = [
    [1, 2],
    [0, 3, 4],
    [0],
    [1],
    [1],
    [6, 7],
    [5],
    [5]
]
print(ob.solve(graph))

入力

graph = [
    [1, 2],
    [0, 3, 4],
    [0],
    [1],
    [1],
    [6, 7],
    [5],
    [5]
]

出力

4

コードのポイント

• dfs1 は各ノードを1度だけ訪問するため、全体の計算量はヒープ操作を含めて O(N log N) 程度に収まります。

• best2 に上位2つの深さのみを保持することで、「あるノードを頂点とした最長パス（左右の部分木の深さの合計）」を効率的に求めることができます。

• 最終的な答えは、「新しく接続してできた木の直径」と「元の森に存在していた最大の直径」のうち大きい方になります。