けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

CS Academy FII Code #1 D - Sugarel in Love

「TDPC うなぎ」の類題。

問題へのリンク

問題概要

 N 頂点の重み付きツリーが与えられる。ツリー上の「頂点を共有しないようなパスの集合」として考えられるもののうち、パスの重みの総和の最大値を求めよ。

制約

  •  2 \le N \le 10^{5}

考えたこと

ツリー二重 DP をする。

ツリー DP のうち、「各ノード v の値を更新すること自体に、v の子供をシーケンシャルに見て行くことによる DP」をするようなものを、個人的にツリー二重 DP と呼んでいる。

  • dp[v][0] := v を根とする部分木において、v を含むパスがない場合の最大値
  • dp[v][1] := v を根とする部分木において、v を含むパスがあって、v がパスの端点になっている場合の最大値
  • dp[v][2] := v を根とする部分木において、v を含むパスがあって、v がパスの内点になっている場合の最大値

として更新すれば OK

#include <iostream>
#include <vector>
using namespace std;
void chmax(long long &a, long long b) { if (a > b) a = b; }

const long long INF = 1LL<<60;
const int MAX = 110000;

int N;
using Edge = pair<int, long long>;
vector<vector<Edge> > G;

long long dp[MAX][3];
void rec(int v, int p = -1) {
    int nch = 0;
    vector<Edge> chs;
    for (auto e : G[v]) {
        if (e.first == p) continue;
        rec(e.first, v);
        ++nch;
        chs.push_back(e);
    }
    vector<vector<long long> > dp2(nch + 1, vector<long long>(3, -INF));
    dp2[0][0] = 0;
    for (int i = 0; i < chs.size(); ++i) {
        Edge e = chs[i];
        int ch = e.first;
        long long add = e.second;

        // j -> j
        for (int j = 0; j < 3; ++j) {
            chmax(dp2[i+1][j], dp2[i][j] + max(dp[ch][0], max(dp[ch][1], dp[ch][2])));
        }
        
        // j -> j+1
        for (int j = 0; j < 2; ++j) {
            chmax(dp2[i+1][j+1], dp2[i][j] + add + max(dp[ch][0], dp[ch][1]));
        }
    }
    for (int j = 0; j < 3; ++j) {
        dp[v][j] = dp2[nch][j];
    }
}

int main() {
    while (cin >> N) {
        G.clear(); G.resize(N);
        for (int i = 0; i < N-1; ++i) {
            int a, b; long long c; cin >> a >> b >> c; --a, --b;
            G[a].emplace_back(b, c);
            G[b].emplace_back(a, c);
        }
        rec(0);
        long long res = 0;
        for (int j = 0; j < 3; ++j) chmax(res, dp[0][j]);
        cout << res << endl;
    }
}