C - 組織のアンバランス / Imbalance of the Organization 解説 by admin
Claude 4.6 Opus (Thinking)概要
根付き木で表される組織において、各社員の直属の部下が率いるチームの仕事量(部分木の総和)の最大値と最小値の差を求め、その全社員にわたる最大値を出力する問題です。
考察
重要な気づき
葉ノードのみが仕事量を持つ: 部下を持つ社員(内部ノード)の \(W_i\) は \(0\) なので、ある社員 \(u\) を根とする部分木の仕事量の総和は、その部分木に含まれる葉ノードの \(W_i\) の合計に等しくなります。
アンバランス度の定義: 社員 \(v\) のアンバランス度は、\(v\) の直属の部下 \(c_1, c_2, \ldots, c_k\) それぞれについて「部分木の仕事量の総和 \(S(c_j)\)」を計算し、その最大値と最小値の差 \(\max(S(c_j)) - \min(S(c_j))\) です。直属の部下が1人以下なら差は取れないので \(0\) です。
やるべきこと: 各ノードの部分木の仕事量の総和を求めれば、あとは各ノードについて子の部分木総和の最大・最小の差を取るだけです。
素朴なアプローチで十分か?
各ノードの部分木の総和は、葉から親に向かって足し上げれば \(O(N)\) で計算できます。その後、各ノードについて子の部分木総和を見て最大・最小を求めるのも、全体で子の数の合計が \(N-1\) なので \(O(N)\) です。したがって、全体 \(O(N)\) で解けるため、\(N \leq 2 \times 10^5\) に対して十分高速です。
アルゴリズム
- 木の構築: 入力から各ノードの子リスト
childrenを作る。 - 部分木の仕事量の総和を計算:
- スタックを使った反復的DFSでトポロジカル順序(根から葉の順)を求める。
- その逆順(葉から根の順)に処理し、葉なら \(S(\text{node}) = W[\text{node}]\)、内部ノードなら \(S(\text{node}) = \sum_{c \in \text{children}(\text{node})} S(c)\) とする。
- アンバランス度の計算:
- 各ノード \(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 によって生成されました。