読者です 読者をやめる 読者になる 読者になる

子から異なる2つを選ぶ感じの木DPを解くアレ

子から異なる2つを選ぶ感じの木DPを解くアレを解くテクニックを紹介します。

D: Longest Path - Indeedなう(オープンコンテスト) | AtCoder

とかが簡単に解けたりします

頂点0を根とする根付き木について

  • 頂点iはa[i]を持つ(a[i] >= 0)
  • 頂点iそれぞれについて、b[i] = a[i] + max(0, max(b[x] ただしxはiの子), max(b[x] + b[y] ただしx, yはどちらもiの子でかつx != y))
  • b[0]を求めよ

という問題を考えます
つまり、b[i]は子供から0, 1, 2個選んだ時のb[選んだ頂点]のsumのmaxにa[i]を足したものですね

vector<int> g[MAX_V]; //g[p]は頂点pの子供を入れたvectorとする
int a[MAX_V], b[MAX_V];
void dfs(int p) {
    for (int x: g[p]) {
        dfs(x);
    }
    int res = a[i];
    for (int x: g[p]) {
        res = max(res, a[i] + b[x]);
    }
    for (int x: g[p]) {
        for (int y: g[p]) {
            if (x == y) continue;
            res = max(res, a[i] + b[x] + b[y]);
        }
    }
    b[i] = res;
    return;
}

これを解く愚直なコードはこんな感じになります。計算量はO(V^2)でオワコンですね。

O(V)にする方法としては
dp[MAX_V][3]を用意して
dp[i][j] : 頂点iについて、子供からj個選んだ時のb[選んだ頂点]の最大値
と置けば

vector<int> g[MAX_V]; //g[p]は頂点pの子供を入れたvectorとする
int a[MAX_V], b[MAX_V];
int dp[MAX_V][3];
void dfs(int p) {
    for (int x: g[p]) {
        dfs(x);
    }
    int res = a[i];
    dp[p][0] = dp[p][1] = dp[p][2] = 0;
    for (int x: g[p]) {
        dp[p][2] = max(dp[p][2], dp[p][1] + b[x]);
        dp[p][1] = max(dp[p][1], dp[p][0] + b[x]);
    }
    b[i] = a[i] + max(max(dp[p][0], dp[p][1]), dp[p][2]);
    return;
}

と解けます
ですが、これが冒頭で紹介した問題のように遷移式が面倒になってくるともうめちゃくちゃですよします

ここで、iの子がu個あるとすると、
(2個選んだ時のmax) = max( a[i] + ([1, j]番目の子のbの最大値 + [j+1, u]番目の子のbの最大値) 1 <= j <= u)
であることに注目します

すると

/**
 * このコードではコードを簡潔にするために
 * 頂点の子を 0,1, .. u-1 番目 と置いて
 * 2個選んだ時のmax = max( [0..i)番目のbのmax + [i.. u)番目のbのmax )
 * と考えています
 */
vector<int> g[MAX_V]; //g[p]は頂点pの子供を入れたvectorとする
int a[MAX_V], b[MAX_V];
void dfs(int p) {
    for (int x: g[p]) {
        dfs(x);
    }
    int u = g[p].size();
    vector<int> left_max(u+1), right_max(u+1);
    //left_max[i]は[0..i)番目の子の中でのbの最大値
    //right_max[i]は[i..u)番目の子の中でのbの最大値
    left_max[0] = 0;
    for (int i = 0; i < u; i++) {
        left_max[i+1] = max(left_max[i], b[g[p][i]]);
    }
    right_max[u] = 0;
    for (int i = u-1; i >= 0; i--) {
        right_max[i] = max(right_max[i+1], b[g[p][i]]);
    }
    int res = a[i]; //0個
    for (int d: g[p]) {
        res = max(res, a[i] + b[d]); //1個(今回の問題では本当はこの更新は不要)
    }
    for (int i = 1; i < u; i++) {
        res = max(res, a[i] + left_max[i] + right_max[i]); //2個
    }
    b[i] = res;
}

となります
コード自体は長くなりますが、バグらせにくい(当社比)

最初の問題も
max ( max( [0 .. i)番目の登るパスの最大長さ + [i .. u)番目の降りるパスの最大長さ)
max( [0 .. i)番目の降りるパスの最大長さ + [i .. u)番目の登るパスの最大長さ) )
を求めればいいため、わりかし簡単になります(当社比)

この挟み込むやつを使うと
max (子から1個だけ選んで、それ以外のsum)
みたいな子を1個えらんでそれ以外にhogeしたもののmaxみたいなのも解けます