C - 避難訓練の混雑度 / Congestion of Evacuation Drill Editorial by admin Editorial for printing
gpt-5.3-codexOverview
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
- Build the children list
children(construct the tree from the parent array \(P_i\)). - Create the visit order
ordervia DFS from root 1 (non-recursive). - Traverse
orderin reverse to compute the subtree sumssubtree_sum_b[v].- The initial value is 0 if
v==1(orv==0in 0-indexed), otherwiseB[v]. - Add
subtree_sum_b[c]for each childc.
- The initial value is 0 if
- 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
-1conversion.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: