概要

根付き木で表される組織において、各社員の直属の部下が率いるチームの仕事量（部分木の総和）の最大値と最小値の差を求め、その全社員にわたる最大値を出力する問題です。

考察

重要な気づき

1. 葉ノードのみが仕事量を持つ: 部下を持つ社員（内部ノード）の \(W_i\) は \(0\) なので、ある社員 \(u\) を根とする部分木の仕事量の総和は、その部分木に含まれる葉ノードの \(W_i\) の合計に等しくなります。

2. アンバランス度の定義: 社員 \(v\) のアンバランス度は、\(v\) の直属の部下 \(c_1, c_2, \ldots, c_k\) それぞれについて「部分木の仕事量の総和 \(S(c_j)\)」を計算し、その最大値と最小値の差 \(\max(S(c_j)) - \min(S(c_j))\) です。直属の部下が1人以下なら差は取れないので \(0\) です。

3. やるべきこと: 各ノードの部分木の仕事量の総和を求めれば、あとは各ノードについて子の部分木総和の最大・最小の差を取るだけです。

素朴なアプローチで十分か？

各ノードの部分木の総和は、葉から親に向かって足し上げれば \(O(N)\) で計算できます。その後、各ノードについて子の部分木総和を見て最大・最小を求めるのも、全体で子の数の合計が \(N-1\) なので \(O(N)\) です。したがって、全体 \(O(N)\) で解けるため、\(N \leq 2 \times 10^5\) に対して十分高速です。

アルゴリズム

1. 木の構築: 入力から各ノードの子リスト children を作る。
2. 部分木の仕事量の総和を計算:
   • スタックを使った反復的DFSでトポロジカル順序（根から葉の順）を求める。
   • その逆順（葉から根の順）に処理し、葉なら \(S(\text{node}) = W[\text{node}]\)、内部ノードなら \(S(\text{node}) = \sum_{c \in \text{children}(\text{node})} S(c)\) とする。
3. アンバランス度の計算:
   • 各ノード \(v\) について、子が2人以上あれば、子の部分木総和 \(S(c)\) の最大値と最小値の差を計算する。
   • すべてのノードにわたってこの差の最大値が答え。

具体例

例えば、社員1が社長で部下が社員2と社員3、社員2の部下が社員4と社員5（いずれも葉）、社員3は葉とします。\(W = [0, 0, 5, 3, 7]\) のとき：

• \(S(4) = 3, S(5) = 7, S(3) = 5\)
• \(S(2) = 3 + 7 = 10\)
• 社員2のアンバランス度: \(7 - 3 = 4\)
• 社員1のアンバランス度: \(\max(10, 5) - \min(10, 5) = 5\)
• 答え: \(\max(4, 5) = 5\)

計算量

• 時間計算量: \(O(N)\)（木の構築、部分木総和の計算、アンバランス度の計算すべて線形）
• 空間計算量: \(O(N)\)（木の隣接リスト、部分木総和の配列、DFS用スタックなど）

実装のポイント

• 再帰ではなく反復DFS: \(N\) が最大 \(2 \times 10^5\) なので、Python の再帰制限（デフォルト1000）に引っかかる可能性があります。スタックを使った反復的なDFSでトポロジカル順序を求め、逆順に処理することで再帰を回避しています。

• 入力の高速読み込み: sys.stdin.buffer.read() で一括読み込みし、split() でトークンに分割することで、大量入力でも高速に処理できます。

• \(N = 1\) の特殊ケース: 社員が1人だけの場合、部下が存在しないためアンバランス度は \(0\) です。親の入力行が空になるため、先に処理しています。

  ソースコード

import sys
from sys import setrecursionlimit

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    
    if N == 1:
        print(0)
        return
    
    children = [[] for _ in range(N + 1)]
    for i in range(2, N + 1):
        p = int(input_data[idx]); idx += 1
        children[p].append(i)
    
    W = [0] * (N + 1)
    for i in range(1, N + 1):
        W[i] = int(input_data[idx]); idx += 1
    
    # Compute subtree sums iteratively
    subtree_sum = [0] * (N + 1)
    
    # Topological order via iterative DFS
    order = []
    stack = [1]
    while stack:
        node = stack.pop()
        order.append(node)
        for c in children[node]:
            stack.append(c)
    
    # Process in reverse order (leaves first)
    for node in reversed(order):
        if not children[node]:
            subtree_sum[node] = W[node]
        else:
            s = 0
            for c in children[node]:
                s += subtree_sum[c]
            subtree_sum[node] = s
    
    # Compute answer
    ans = 0
    for v in range(1, N + 1):
        if len(children[v]) >= 2:
            child_sums = [subtree_sum[c] for c in children[v]]
            max_val = max(child_sums)
            min_val = min(child_sums)
            diff = max_val - min_val
            if diff > ans:
                ans = diff
    
    print(ans)

main()

この解説は claude4.6opus-thinking によって生成されました。