公式
C - 避難訓練の混雑度 / Congestion of Evacuation Drill 解説
by
C - 避難訓練の混雑度 / Congestion of Evacuation Drill 解説
by
MtSaka
各 \(i\) について \(C_i=\sum_{j\in \text{path}(i,1)}A_j\) を求め、この \(C\) について \(\sum_{i=2}^{N}B_iC_i\) が答えとなります。
広場 \(i\) から広場 \(1\) までの単純パスにおいては広場 \(i\) に隣接するパス上の点として広場 \(P_i\) があります。 つまり、\(\sum_{j\in \text{path}(i,1)}A_j=A_i+\sum_{j\in \text{path}(P_i,1)}A_j\) が成り立ちます。この式を用いて、\(C_1=A_1\) として、 \(i=2,3,\ldots,\) と昇順に\(C_i\) を計算することができます。昇順に計算してよいのは、\(P_i<i\) が成り立つためです。
したがって、全体で \(\mathrm{O}(N)\) 時間で解くことができます。
実装例(C++)
#include <bits/stdc++.h>
using namespace std;
int main() {
int n;
cin >> n;
vector<int> a(n);
for (auto& e : a) cin >> e;
vector<int> p(n, -1);
for (int i = 1; i < n; ++i) cin >> p[i];
for (auto& e : p) e--;
vector<int> b(n);
for (auto& e : b) cin >> e;
vector<long long> c(n);
c[0] = a[0];
for (int i = 1; i < n; ++i) c[i] = c[p[i]] + a[i];
long long ans = 0;
for (int i = 1; i < n; ++i) ans += b[i] * c[i];
cout << ans << endl;
}
