C - 避難訓練の混雑度 / Congestion of Evacuation Drill 解説 by admin
GPT 5.4 HighOverview
This problem asks us to sum, for each person, “the total of \(A\) values of the plazas on the path from their current plaza to the root \(1\).”
By changing the perspective to “how many people pass through each plaza,” this can be computed in \(O(N)\) as a subtree sum on the tree.
Analysis
A straightforward approach would be, for the \(B_i\) people at each plaza \(i\), to follow the parent pointers from \(i\) to the root \(1\), accumulating the congestion coefficients along the way.
However, this method processes a number of nodes equal to the distance from each vertex to the root.
For example, if the tree is a straight line, vertex \(N\) requires \(N\) steps, vertex \(N-1\) requires \(N-1\) steps, and so on, leading to an overall complexity of \(O(N^2)\).
Since \(N \le 3 \times 10^5\), this is too slow.
Let us reverse the order of summation.
The formula in the problem is
\[ \sum_{i=2}^{N} B_i \sum_{j \in \mathrm{path}(i,1)} A_j \]
This represents the perspective:
- “For each person, sum the \(A\) values of the plazas they pass through.”
We can reinterpret this as:
- “For each plaza \(j\), add (the number of people passing through that plaza) × \(A_j\).”
Then the answer becomes
\[ \sum_{j=1}^{N} A_j \times (\text{number of people passing through plaza } j) \]
So, how many people pass through plaza \(j\)?
- When \(j \ne 1\):
The people who pass through plaza \(j\) are exactly “the people in the subtree of \(j\).”
This is because, no matter which vertex in the subtree you start from, the path to the root necessarily passes through \(j\).
Therefore,
\[ \text{number of people passing through plaza } j = \text{total number of people in subtree } j \]
- However, we need to be careful when \(j = 1\).
The \(B_1\) people at the root \(1\) “do not move,” so according to the problem statement, they do not pass through any plaza.
Thus, the number of people passing through plaza \(1\) is
\[ \text{total number of people in subtree } 1 - B_1 \]
So, if we compute for each vertex \(i\)
\[ \text{sub}[i] = \text{total number of people in the subtree of vertex } i \]
then the answer is
\[ A_1 \times (\text{sub}[1] - B_1) + \sum_{i=2}^{N} A_i \times \text{sub}[i] \]
How to compute subtree sums
Normally, we would use DFS to compute subtree sums, but in this problem we can do it more simply.
The constraint states
\[ P_i < i \]
so a parent’s index is always smaller than its child’s.
This means that by iterating from the largest vertex index downward, we can directly accumulate each child’s information into its parent.
Initialize
\[ \text{sub}[i] = B_i \]
then, for \(i = N, N-1, \dots, 2\) in order, perform
\[ \text{sub}[P_i] += \text{sub}[i] \]
This gives us the subtree population for each vertex.
This is a very important point — there is no need to explicitly build an adjacency list for the tree, making the implementation very concise.
Concrete example
Consider the following tree:
- \(1\) is the root
- The parent of \(2, 3\) is \(1\)
- The parent of \(4, 5\) is \(2\)
In this case, the people passing through plaza \(2\) are:
- People at \(2\)
- People at \(4\)
- People at \(5\)
In other words, “everyone in the subtree of \(2\).”
Similarly, the people passing through plaza \(1\) are everyone in the entire tree, but the people originally at \(1\) do not move, so they are excluded.
This perspective — “the number of people passing through each vertex = the subtree population” — is the essence of this problem.
Algorithm
- Initialize the array \(\text{sub}\) with \(\text{sub}[i] = B_i\).
- Iterate in the order \(i = N, N-1, \dots, 2\), performing \(\text{sub}[P_i] += \text{sub}[i]\).
This computes the subtree population for each vertex. - Compute the answer as follows:
\[ \text{ans} = A_1 \times (\text{sub}[1] - B_1) \]
\[ \text{ans} += \sum_{i=2}^{N} A_i \times \text{sub}[i] \]
- Output \(\text{ans}\).
Complexity
- Time complexity: \(O(N)\)
- Space complexity: \(O(N)\)
Implementation notes
Note that the root \(1\) is special: the number of people passing through it is \(\text{sub}[1] - B_1\), not \(\text{sub}[1]\).
Because this problem guarantees \(P_i < i\), we can compute subtree sums with a single backward pass without using DFS.
Using 1-indexed arrays in the implementation makes it easier to correspond with the formulas and reduces bugs.
In Python, integers automatically support arbitrary precision, so there is no need to worry about overflow.
Source Code
import sys
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
it = iter(data)
N = next(it)
A = [0] + [next(it) for _ in range(N)]
P = [0] * (N + 1)
for i in range(2, N + 1):
P[i] = next(it)
B = [0] + [next(it) for _ in range(N)]
sub = B[:]
for i in range(N, 1, -1):
sub[P[i]] += sub[i]
ans = A[1] * (sub[1] - B[1])
for i in range(2, N + 1):
ans += A[i] * sub[i]
print(ans)
if __name__ == "__main__":
main()
This editorial was generated by gpt-5.4-high.