Official

C - 避難訓練の混雑度 / Congestion of Evacuation Drill Editorial by admin Editorial for printing

gpt-5.3-codex

Overview

Rather than tracing the path for each person and summing up individually, this problem asks us to count “how many people’s totals does each plaza’s congestion coefficient \(A_v\) ultimately contribute to.”
By using the sum of people in each subtree of the tree, we can compute the entire answer in \(O(N)\).

Analysis

The formula is \( \displaystyle \sum_{i=2}^{N} B_i \sum_{j \in \mathrm{path}(i,1)} A_j \) If we naively “climb from each \(i\) to the root,” in the worst case we traverse a path of length \(O(N)\) from each vertex, resulting in \(O(N^2)\) time, which is too slow for \(N=3\times 10^5\).

So we change the perspective of the summation.

  • Instead of counting from the person (starting point) side,
  • Consider from the plaza \(v\) side: “how many people is this \(A_v\) added for?”

The people who pass through plaza \(v\) are “those in the subtree of \(v\) (excluding the person at plaza 1, who does not move).”
In other words, the contribution of \(A_v\) is \( A_v \times (\text{number of movers in the subtree of } v) \)

Therefore, the answer is \( \displaystyle \sum_{v=1}^{N} A_v \cdot S_v \) where \(S_v\) is “the sum of \(B\) values within the subtree of \(v\) (excluding vertex 1 itself).”

All that remains is to perform tree DP (subtree sums) once.

Algorithm

  1. Build the children list children (construct the tree from the parent array \(P_i\)).
  2. Create the visit order order via DFS from root 1 (non-recursive).
  3. Traverse order in reverse to compute the subtree sums subtree_sum_b[v].
    • The initial value is 0 if v==1 (or v==0 in 0-indexed), otherwise B[v].
    • Add subtree_sum_b[c] for each child c.
  4. Finally, compute and output \( \sum_v A[v]\times subtree\_sum\_b[v] \)

In this implementation, the condition that the people at the root \(B_1\) do not move is naturally handled by setting the initial value to 0 when v==0.

Complexity

  • Time complexity: \(O(N)\)
  • Space complexity: \(O(N)\)

Implementation Notes

  • The input is 1-indexed and Python arrays are 0-indexed, so don’t forget the -1 conversion.

  • Recursive DFS may hit the recursion limit on deep trees, so we use an iterative DFS with a stack instead.

  • The answer can be very large, but Python’s integers support arbitrary precision, so it is safe to use as-is (the problem statement also guarantees it fits within 64 bits).

    Source Code

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()

This editorial was generated by gpt-5.3-codex.

posted:
last update: