公式

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

GPT 5.4 High

概要

各人について「自分のいる広場から根 \(1\) までの経路上の広場の \(A\) の総和」を足す問題です。
これを「各広場が、何人に通過されるか」を数える形に見方を変えると、木の部分木和として \(O(N)\) で求められます。

考察

素朴に考えると、各広場 \(i\) にいる \(B_i\) 人について、\(i\) から根 \(1\) まで親をたどりながら混雑係数を足せばよさそうです。

しかしこの方法では、各頂点から根までの長さぶん処理することになります。
たとえば木が一直線だと、頂点 \(N\) からは \(N\) 個、頂点 \(N-1\) からは \(N-1\) 個……となり、全体で \(O(N^2)\) かかってしまいます。
\(N \le 3 \times 10^5\) なので、これは間に合いません。


ここで和の取り方を逆にしてみます。

問題の式は

\[ \sum_{i=2}^{N} B_i \sum_{j \in \mathrm{path}(i,1)} A_j \]

です。
これは

  • 「各人ごとに、通る広場の \(A\) を足す」

という見方ですが、これを

  • 「各広場 \(j\) について、その広場を通る人数 × \(A_j\) を足す」

と考え直せます。

すると答えは

\[ \sum_{j=1}^{N} A_j \times (\text{広場 } j \text{ を通る人数}) \]

になります。


では、広場 \(j\) を通る人数は何人でしょうか?

  • \(j \ne 1\) のとき
    広場 \(j\) を通るのは、「\(j\) の部分木内にいる人たち」です。
    なぜなら、部分木内のどの頂点から根へ向かっても、必ず \(j\) を通るからです。

よって、

\[ \text{広場 } j \text{ を通る人数} = \text{部分木 } j \text{ にいる人数の総和} \]

です。

  • ただし \(j = 1\) のときだけ注意が必要です。
    根 \(1\) にいる \(B_1\) 人は「移動しない」ので、問題文よりどの広場も通過しません。
    したがって、広場 \(1\) を通る人数は

\[ \text{部分木 } 1 \text{ にいる人数の総和} - B_1 \]

となります。


そこで、各頂点 \(i\) について

\[ \text{sub}[i] = \text{頂点 } i \text{ の部分木にいる人数の総和} \]

を求めれば、答えは

\[ A_1 \times (\text{sub}[1] - B_1) + \sum_{i=2}^{N} A_i \times \text{sub}[i] \]

となります。


部分木和はどう求めるか

普通は DFS を使って部分木和を求めますが、この問題ではもっと簡単にできます。

制約に

\[ P_i < i \]

とあるので、親の番号は必ず子より小さいです。
つまり、頂点番号の大きい方から見ていけば、子の情報を親へそのまま足し上げられます。

初めに

\[ \text{sub}[i] = B_i \]

としておき、\(i=N,N-1,\dots,2\) の順に

\[ \text{sub}[P_i] += \text{sub}[i] \]

とすれば、各頂点の部分木人数が求まります。

これはとても重要なポイントで、木をわざわざ隣接リストにしなくてもよいので実装も簡潔です。


具体例

たとえば

  • \(1\) が根
  • \(2,3\) の親は \(1\)
  • \(4,5\) の親は \(2\)

という木を考えます。

このとき、広場 \(2\) を通るのは

  • \(2\) にいる人
  • \(4\) にいる人
  • \(5\) にいる人

です。
つまり「\(2\) の部分木にいる全員」です。

同様に、広場 \(1\) を通るのは木全体にいる人ですが、\(1\) に最初からいる人は移動しないので除外します。

この「各頂点を通る人数 = 部分木人数」という見方が、この問題の本質です。

アルゴリズム

  1. 配列 \(\text{sub}\) を \(\text{sub}[i] = B_i\) で初期化する。
  2. \(i = N, N-1, \dots, 2\) の順に見て、\(\text{sub}[P_i] += \text{sub}[i]\) を行う。
    これで各頂点の部分木内人数が求まる。
  3. 答えを次で計算する。

\[ \text{ans} = A_1 \times (\text{sub}[1] - B_1) \]

\[ \text{ans} += \sum_{i=2}^{N} A_i \times \text{sub}[i] \]

  1. \(\text{ans}\) を出力する。

計算量

  • 時間計算量: \(O(N)\)
  • 空間計算量: \(O(N)\)

実装のポイント

  • 根 \(1\) だけは特別で、通る人数が \(\text{sub}[1]\) ではなく \(\text{sub}[1] - B_1\) になる点に注意しましょう。

  • この問題では \(P_i < i\) があるので、DFS を使わずに後ろから 1 回なめるだけで部分木和を計算できます。

  • 実装では 1-indexed にすると式と対応しやすく、バグを減らせます。

  • Python では整数が自動で多倍長になるためオーバーフローは心配ありません。

    ソースコード

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

この解説は gpt-5.4-high によって生成されました。