Pythonで木の辺を1本取り除いたときの部分木のノード値合計の差の最小値を求めるプログラム

問題の概要

ノードに1からnまでの番号が振られた木があるとします。各ノードには整数値が格納されています。ここで、木からある1本の辺を取り除くと、木は2つの部分木に分割されます。このとき、2つの部分木のノード値の合計の差が最小になるようにしたいと考えます。私たちのタスクは、その最小の差を求めて返すことです。木は辺のリストとして与えられ、各ノードの値も併せて提供されます。

例として、n = 6、edge_list = [[1, 2], [1, 3], [2, 4], [3, 5], [3, 6]]、values = [15, 25, 15, 55, 15, 65] が入力された場合、出力は 0 になります。

この入力に対して、各辺を取り除いた場合の結果を確認してみましょう。

• 辺 (1,2) を取り除くと、重みの合計は 80 と 110 になり、差は 30 です。
• 辺 (1,3) を取り除くと、重みの合計は 95 と 95 になり、差は 0 です。
• 辺 (2,4) を取り除くと、重みの合計は 55 と 135 になり、差は 80 です。
• 辺 (3,5) を取り除くと、重みの合計は 15 と 175 になり、差は 160 です。
• 辺 (3,6) を取り除くと、重みの合計は 65 と 125 になり、差は 60 です。

したがって、最小の差は 0 となります。

解法のアプローチ

この問題を解くために、以下の手順に従います。

• サイズ n の隣接リスト(adj_list)を作成し、各要素を空のリストで初期化します。
• edge_list の各辺について、次の処理を行います。
  • u := edge[0]、v := edge[1] とします。
  • adj_list[u-1] の末尾に v-1 を追加します。
  • adj_list[v-1] の末尾に u-1 を追加します。
• サイズ n の value_list を 0 で初期化して作成します。
• not_visited を、隣接リストの要素数が1である(葉に相当する)ノードの集合として初期化します。
• not_visited が空になるまで、次の処理を繰り返します。
  • not_visited 内の各 i について、value_list[i] に values[i] を加算し、隣接ノードが存在する場合はその親側ノードから自分を削除し、親側ノードの value_list に自分の value_list[i] を加算します(葉から順に部分木の合計を積み上げていくイメージです)。
  • 次の not_visited を、処理後に隣接リストの長さが1になったノードの集合として更新します。
• return_val := |sum(values) - 2 * value_list[0]| とします。
• i が 1 から n-1 までの範囲で、decision_val := |sum(values) - 2 * value_list[i]| を計算し、return_val より小さければ return_val を更新します。
• return_val を返します。

ポイントは、「全体の合計 - 2 × 部分木の合計」の絶対値を取ることで、片方の部分木の合計だけから差を求められる点です。これにより、毎回両方の部分木の合計を計算する必要がなくなります。

実装例

理解を深めるために、以下の実装を見てみましょう。

def solve(n, edge_list, values):
    adj_list = [[] for i in range(n)]

    for edge in edge_list:
        u = edge[0]
        v = edge[1]
        adj_list[u-1].append(v-1)
        adj_list[v-1].append(u-1)

    value_list = [0] * n
    not_visited = {i for i in range(n) if len(adj_list[i]) == 1}
    while(len(not_visited)):
        for i in not_visited:
            value_list[i] += values[i]
            if(len(adj_list[i])):
                adj_list[adj_list[i][0]].remove(i)
                value_list[adj_list[i][0]] += value_list[i]
        not_visited = {adj_list[i][0] for i in not_visited if
            len(adj_list[i]) and len(adj_list[adj_list[i][0]]) == 1}

    return_val = abs(sum(values) - 2 * value_list[0])

    for i in range(1, n):
        decision_val = abs(sum(values) - 2 * value_list[i])
        if decision_val < return_val:
            return_val = decision_val
    return return_val

print(solve(6, [[1, 2], [1, 3], [2, 4], [3, 5], [3, 6]], [10, 20, 10, 50, 10, 60]))

入力

6, [[1, 2], [1, 3], [2, 4], [3, 5], [3, 6]], [10, 20, 10, 50, 10, 60]

出力

0

まとめ

このアルゴリズムでは、葉ノードから順に部分木の合計値を累積していくことで、すべてのノードを一度ずつ処理するだけで各ノードを根とする部分木の合計を求められます。その後、全ノードについて「全体の合計 − 2 × 部分木の合計」の絶対値を評価し、その最小値を返すことで、辺を1本取り除いたときの2つの部分木の合計差の最小値が得られます。計算量は O(n) 程度に抑えられるため、大きな木に対しても効率的に動作します。