全方位木DPメモ

全方位木DPの実装を忘れたときのためのメモ

問題

• http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_5_A&lang=jp

解答

NOTE: AOJはC++17に対応してないためC++14で書いてる。

int d[100010];

void solve() {
    int N;
    cin >> N;
    using Graph = vector<vector<Pii>>;
    Graph g(N);
    for (ll i = 0; i < N - 1; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        g[a].emplace_back(b, c);
        g[b].emplace_back(a, c);
    }

    auto dfs1 = [&](auto&& self, int idx, int par) -> int {
        int res = 0;
        for (auto edge : g[idx]) {
            int u, c;
            tie(u, c) = edge;
            if (u == par) continue;
            chmax(res, self(self, u, idx) + c);
        }
        d[idx] = res;
        return res;
    };
    dfs1(dfs1, 0, -1);

    int ans = 0;
    auto dfs2 = [&](auto&& self, int idx, int par, int pval) -> void {
        int nc = (int)g[idx].size();
        Vi vec(nc);
        Vi lcum(nc + 1), rcum(nc + 1);
        for (int i = 0; i < nc; i++) {
            int u, c;
            tie(u, c) = g[idx][i];
            int t;
            if (u == par) {
                t = pval + c;
            }
            else {
                t = d[u] + c;
            }
            vec[i] = t;
        }
        for (int i = 0; i < nc; i++) {
            lcum[i + 1] = max(lcum[i], vec[i]);
        }
        for (int i = nc - 1; i >= 0; i--) {
            rcum[i] = max(rcum[i + 1], vec[i]);
        }
        sort(vec.begin(), vec.end(), greater<int>());
        if (nc == 1) {
            chmax(ans, vec[0]);
        }
        else if (nc > 1) {
            chmax(ans, vec[0] + vec[1]);
        }
        for (int i = 0; i < nc; i++) {
            int u, c;
            tie(u, c) = g[idx][i];
            if (u == par) continue;
            self(self, u, idx, max(lcum[i], rcum[i + 1]));
        }
    };
    dfs2(dfs2, 0, -1, 0);

    prints()(ans);
}