Official

E - Pair Annihilation Editorial by sotanishy


与えられた木を,頂点 \(1\) を根とする根付き木とみなします.また,陽電子と電子をまとめて「粒子」と呼ぶことにします.

\(j\) が答えにどれくらい寄与するかを考えます.一般性を失わず,頂点 \(u_j\)\(v_j\) の親であるとします.

頂点 \(v_j\) を根とする部分木の頂点に関する \(x\) の総和を \(X_j\) とします.この部分木の中だけでどのように操作を行っても,\(|X_j|\) 個の粒子が残ります.残った粒子は,辺 \(j\) を通して部分木の外に移動させるか,または辺 \(j\) を通して部分木に入ってきた粒子と打ち消しあう必要があります.どちらにせよ,少なくとも \(|X_j|\) 個の粒子を辺 \(j\) に沿って移動させる必要があり,答えに少なくとも \(|X_j| w_j\) だけ寄与します.よって,答えは少なくとも \(\sum_{j=1}^{N-1} |X_j| w_j\) です.

そして,この下界を達成する操作が存在します.具体的には,葉に近い頂点から見ていって,その頂点に粒子が残っているならば,親に向かって移動させるということを繰り返せばよいです.

よって最適解は \(\sum_{j=1}^{N-1} |X_j| w_j\) です.\(X_j\)深さ優先探索 (DFS) により求めることができます.時間計算量は \(O(N)\) です.

実装例 (Python)

N = int(input())
x = list(map(int, input().split()))
T = [[] for _ in range(N)]
for _ in range(N - 1):
    u, v, w = map(int, input().split())
    u -= 1
    v -= 1
    T[u].append((v, w))
    T[v].append((u, w))

ans = 0
st = [(0, -1, 0)]
while st:
    v, p, t = st.pop()
    if t == 0:
        st.append((v, p, 1))
        for u, w in T[v]:
            if u != p:
                st.append((u, v, 0))
    else:
        for u, w in T[v]:
            if u == p:
                continue
            ans += w * abs(x[u])
            x[v] += x[u]
print(ans)

実装例 (C++)

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N;
    cin >> N;
    vector<long long> x(N);
    for (auto& y : x) cin >> y;
    vector<vector<pair<int, long long>>> T(N);
    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        long long w;
        cin >> u >> v >> w;
        --u, --v;
        T[u].push_back({v, w});
        T[v].push_back({u, w});
    }

    long long ans = 0;

    auto dfs = [&](auto& dfs, int v, int p) -> void {
        for (auto [c, w] : T[v]) {
            if (c == p) continue;
            dfs(dfs, c, v);
            ans += w * abs(x[c]);
            x[v] += x[c];
        }
    };

    dfs(dfs, 0, -1);
    cout << ans << endl;
}

posted:
last update: