C - 避難訓練の混雑度 / Congestion of Evacuation Drill Editorial by admin Editorial for printing
gpt-5.3-codex概要
各人ごとに経路をたどって足し上げるのではなく、「各広場の混雑係数 \(A_v\) が最終的に何人分カウントされるか」を数える問題です。
木の部分木にいる人数の合計を使うと、全体を \(O(N)\) で計算できます。
考察
式は
\(
\displaystyle \sum_{i=2}^{N} B_i \sum_{j \in \mathrm{path}(i,1)} A_j
\)
です。
このまま「各 \(i\) について根まで登る」をすると、最悪で各頂点から長さ \(O(N)\) の経路をたどるため \(O(N^2)\) になり、\(N=3\times 10^5\) では間に合いません。
そこで和の見方を変えます。
- 人(始点)側で数えるのでなく、
- 広場 \(v\) 側で「この \(A_v\) は何人に対して加算されるか」を考える。
広場 \(v\) を通るのは、「\(v\) の部分木にいる人(ただし広場1の人は移動しない)」です。
つまり、\(A_v\) の寄与は
\(
A_v \times (\text{\)v\( の部分木にいる移動者数})
\)
になります。
よって答えは \( \displaystyle \sum_{v=1}^{N} A_v \cdot S_v \) ただし \(S_v\) は「\(v\) の部分木内の \(B\) の総和(ただし頂点1自身は除外)」です。
あとは木DP(部分木和)を1回やればよいです。
アルゴリズム
- 子リスト
childrenを作る(親配列 \(P_i\) から木を構築)。 - 根 1 から DFS(再帰なし)で訪問順
orderを作る。 orderを逆順に見て部分木和subtree_sum_b[v]を計算する。- 初期値は
v==1(0-indexならv==0)なら 0、そうでなければB[v] - 子
cのsubtree_sum_b[c]を足す
- 初期値は
- 最後に \( \sum_v A[v]\times subtree\_sum\_b[v] \) を計算して出力。
この実装では、根の人数 \(B_1\) は移動しない条件を
v==0 のとき初期値を 0 にすることで自然に処理しています。
計算量
- 時間計算量: \(O(N)\)
- 空間計算量: \(O(N)\)
実装のポイント
入力は 1-indexed、Python 配列は 0-indexed なので
-1変換を忘れない。再帰DFSは深い木で再帰上限に当たる可能性があるため、スタックで反復DFSにしている。
答えは大きくなるが、Python の整数は多倍長なのでそのままで安全(問題文でも 64bit 以内保証)。
ソースコード
import sys
def main():
input = sys.stdin.readline
N = int(input().strip())
A = list(map(int, input().split()))
P = list(map(int, input().split()))
B = list(map(int, input().split()))
children = [[] for _ in range(N)]
for i in range(2, N + 1):
p = P[i - 2] - 1
children[p].append(i - 1)
# subtree_sum_b[v] = sum of B_u for all u in subtree(v), excluding node 1 itself
subtree_sum_b = [0] * N
order = []
stack = [0]
while stack:
v = stack.pop()
order.append(v)
for c in children[v]:
stack.append(c)
for v in reversed(order):
s = 0 if v == 0 else B[v]
for c in children[v]:
s += subtree_sum_b[c]
subtree_sum_b[v] = s
ans = 0
for v in range(N):
ans += A[v] * subtree_sum_b[v]
print(ans)
if __name__ == "__main__":
main()
この解説は gpt-5.3-codex によって生成されました。
posted:
last update: