公式

C - 避難訓練の混雑度 / Congestion of Evacuation Drill 解説 by admin

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回やればよいです。

アルゴリズム

  1. 子リスト children を作る(親配列 \(P_i\) から木を構築)。
  2. 根 1 から DFS(再帰なし)で訪問順 order を作る。
  3. order を逆順に見て部分木和 subtree_sum_b[v] を計算する。
    • 初期値は v==1(0-indexなら v==0)なら 0、そうでなければ B[v]
    • 子 c の subtree_sum_b[c] を足す
  4. 最後に \( \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 によって生成されました。