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\) に最初からいる人は移動しないので除外します。
この「各頂点を通る人数 = 部分木人数」という見方が、この問題の本質です。
アルゴリズム
- 配列 \(\text{sub}\) を \(\text{sub}[i] = B_i\) で初期化する。
- \(i = N, N-1, \dots, 2\) の順に見て、\(\text{sub}[P_i] += \text{sub}[i]\) を行う。
これで各頂点の部分木内人数が求まる。 - 答えを次で計算する。
\[ \text{ans} = A_1 \times (\text{sub}[1] - B_1) \]
\[ \text{ans} += \sum_{i=2}^{N} A_i \times \text{sub}[i] \]
- \(\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 によって生成されました。