Official

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

gemini-3-flash-thinking

Overview

This problem asks us to find, in a rooted tree, the sum of vertex weights (congestion coefficients) along the path from each vertex to the root (plaza 1), multiply each by the number of people at that vertex, and compute the total sum.

Analysis

Organizing the formula

The “congestion level” we want to compute is expressed by the following formula: $\( \text{Congestion Level} = \sum_{i=2}^{N} (B_i \times \text{sum of } A_j \text{ on the path from plaza } i \text{ to plaza } 1) \)$

If we define \(S_i\) as “the sum of congestion coefficients \(A_j\) of all plazas on the path from plaza \(i\) to plaza \(1\),” then the desired value becomes \(\sum_{i=2}^{N} B_i S_i\), a very simple form.

Naive approach and its limitations

For each plaza \(i\), we could follow parent pointers up to plaza \(1\), adding \(A_j\) along the way. However, if the tree has a long “linear” structure, computing a single path takes \(O(N)\), resulting in \(O(N^2)\) total computation. Under the constraint \(N = 3 \times 10^5\), this will not finish within the time limit.

Efficient computation method

By leveraging the properties of tree structures, \(S_i\) can be computed efficiently using parent information. Let \(P_i\) be the parent of plaza \(i\). Then the following relation holds: $\( S_i = S_{P_i} + A_i \)\( In other words, "the sum from plaza \)i$ to the root” equals “the sum from the parent to the root” plus “the value of the vertex itself.”

In this problem, the constraint \(1 \leq P_i < i\) is given. This means that if we compute in increasing order of index, the computation for parent \(P_i\) is always completed before we compute plaza \(i\). Therefore, we can compute all \(S_i\) in \(O(N)\) with just a simple loop, without needing to implement complex traversals (such as DFS).

Algorithm

  1. Prepare an array \(S\) to store the path sum for each plaza \(i\).
  2. Set the path sum of the root (plaza \(1\)) to \(S_1 = A_1\).
  3. For \(i = 2, 3, \dots, N\) in order, perform the following computation:
    • \(S_i = S_{P_i} + A_i\)
  4. Compute \(B_i \times S_i\) for all \(i = 2, 3, \dots, N\), and output the total sum.

Complexity

  • Time complexity: \(O(N)\)
    • Reading input, computing path sums (one loop), and computing the final total sum (one loop) are all linear in \(N\).
  • Space complexity: \(O(N)\)
    • \(O(N)\) memory is used for arrays storing \(A, P, B, S\).

Implementation Notes

  • Fast input: When handling input of size \(N = 3 \times 10^5\) in Python, repeatedly calling input() is slow, so the standard technique is to read everything at once using sys.stdin.read().split().

  • Index offset: The problem uses 1-indexed notation (plazas 1 through N), but in the program we often use 0-indexed notation (list indices 0 through N-1), so adjustments such as subtracting 1 from parent index \(P_i\) are necessary.

  • Large numbers: The answer can be as large as \(2^{63}-1\) (approximately the maximum value of a 64-bit integer), but since Python natively supports arbitrary-precision integers, there is no concern about overflow.

    Source Code

import sys

def solve():
    # 入力を一括で読み込み、スペース区切りで分割します
    # N=3*10^5 という制約に対して高速に動作させるための手法です
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # 広場の数 N を取得
    N = int(input_data[0])
    
    # A: 各広場の混雑係数 A_1, A_2, ..., A_N
    # リストのインデックス 0 が広場 1 に対応するようにします
    A = list(map(int, input_data[1:N+1]))
    
    # P: 広場 i (2 <= i <= N) の親広場 P_i
    # P[0] は広場 2 の親、P[1] は広場 3 の親... という順で格納されます
    P = list(map(int, input_data[N+1:2*N]))
    
    # B: 各広場にいる人数 B_1, B_2, ..., B_N
    B = list(map(int, input_data[2*N:3*N]))
    
    # S[i]: 広場 i+1 から広場 1 までのパス上の混雑係数の総和
    # 動的計画法(または木上の計算)を用いて各広場からの累積和を求めます
    S = [0] * N
    S[0] = A[0]  # 広場 1 のパス和は A_1 自体
    
    # 制約 1 <= P_i < i により、親のパス和が常に先に計算されていることが保証されます
    # これにより、トポロジカル順序(この場合は単なるインデックス順)で計算可能です
    for i in range(1, N):
        # i は 0-indexed なので、広場番号としては i+1
        # その親の広場番号は P[i-1]
        parent_idx = P[i-1] - 1
        # 広場 i+1 のパス和 = 広場 i+1 の混雑係数 + 親広場のパス和
        S[i] = S[parent_idx] + A[i]
        
    # 全体の混雑度を計算します
    # 広場 1 にいる B_1 人は移動しないため、i=2 から N までの和をとります
    # (0-indexed では i=1 から N-1 まで)
    ans = 0
    for i in range(1, N):
        ans += B[i] * S[i]
        
    # 結果を出力
    sys.stdout.write(str(ans) + '\n')

if __name__ == '__main__':
    solve()

This editorial was generated by gemini-3-flash-thinking.

posted:
last update: