Typical DP Contest P - うなぎ

競プロ TDPC AtCoder

問題文

http://tdpc.contest.atcoder.jp/tasks/tdpc_eel

解法

木の上でのDPなので，やはり根を適当に決めたあとのDFS木上で，部分木の結果を統合していくのが鉄則となります．

根として適当に頂点 0 を選んで，そのDFS木上で考えます．
以下のような dp を考えます．
dp[v][i][j] := 根が v となるような部分木で，i 個の disjoint なパスを書き，かつ v の次数が j であるような書き方

パスは disjoint であるという条件から，k はたかだか 2 になります．
dp の更新にさらに別のDPを行う必要があります．
dp2[i][j][k] := v の子のうち i 個目まで見た時，j 個の disjoint なパスを書き，かつ v の次数が k であるような書き方
dp2 の更新には，子の dp の値も使います．

dp2 の遷移は，基本的に子 i と v に辺を書くかどうかで決まります．
以下の4パターンがあります．

  • そもそも書かない場合 -> 積を単純に加えるだけでよい．
  • 書く場合
    • v の次数が 0 で，i の次数が 0 -> v と i だけからなるパスが新たに1つ増える
    • v の次数が 0 で，i の次数が 1 -> もともとあるパスが延長される
    • v の次数が 1 で，i の次数が 0 -> もともとあるパスが延長される
    • v の次数が 1 で，i の次数が 1 -> もともと disjoint だったパス同士がくっつくので，トータルで disjoint なパスが1つ減る

これを丁寧に書くと良いです．
ちょっと説明が雑なので，コードを読んだほうがわかるかもしれません．

ソースコード

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

using ll = long long;
using graph = vector<vector<int>>;

constexpr ll M = 1e9 + 7;

void dfs(graph const& g, int v, int par, vector<vector<vector<ll>>>& dp) {
    vector<int> child;
    for(auto to : g[v]) {
        if(to == par) {
            continue;
        }
        child.push_back(to);
        dfs(g, to, v, dp);
    }
    int const sz = child.size();
    int const N = dp.size();
    int const K = dp[0].size() - 1;

    vector<vector<vector<ll>>> dp2(sz + 1, vector<vector<ll>>(K + 1, vector<ll>(3)));
    dp2[0][0][0] = 1;
    // i: 今見ている子の idx
    // j: 頂点 v の部分木のパスの数
    // k: 頂点 v の次数
    // l: 今見ている子の部分木でのパスの数
    // m: 今見ている子の部分木での子の次数
    for(int i = 0; i < sz; ++i) {
        for(int j = 0; j < K + 1; ++j) {
            for(int k = 0; k < 3; ++k) {
                for(int l = 0; l < K + 1; ++l) {
                    for(int m = 0; m < 3; ++m) {
                        // パスが一個減るかもしれないので K + 1 までは許容
                        if(j + l > K + 1) {
                            continue;
                        }

                        ll t = (dp2[i][j][k] * dp[child[i]][l][m]) % M;
                        // v と child[i] を結ばない
                        if(j + l <= K) {
                            (dp2[i + 1][j + l][k] += t) %= M;
                        }

                        // v と child[i] を結ぶ
                        if(k == 0) {
                            if(m == 0 && j + l + 1 <= K) {
                                // v と child[i] を結ぶことで v - child[i] のみの新たなパスが増える
                                (dp2[i + 1][j + l + 1][1] += t) %= M;
                            } else if(m == 1 && j + l <= K) {
                                (dp2[i + 1][j + l][1] += t) %= M;
                            }
                        } else if(k == 1) {
                            if(m == 0 && j + l <= K) {
                                (dp2[i + 1][j + l][2] += t) %= M;
                            } else if(m == 1 && 0 <= j + l - 1) {
                                // 2つのパスがくっつくので 1 つ減る
                                (dp2[i + 1][j + l - 1][2] += t) %= M;
                            }
                        }
                    }
                }
            }
        }
    }
    dp[v] = move(dp2[sz]);
}

int main() {
    int N, K;
    cin >> N >> K;
    graph g(N);
    for(int i = 0; i < N - 1; ++i) {
        int a, b;
        cin >> a >> b;
        a--; b--;
        g[a].push_back(b);
        g[b].push_back(a);
    }

    vector<vector<vector<ll>>> dp(N + 1, vector<vector<ll>>(K + 1, vector<ll>(3)));
    dfs(g, 0, -1, dp);

    ll res = 0;
    for(int i = 0; i < 3; ++i) {
        (res += dp[0][K][i]) %= M;
    }

    cout << res << endl;
}

感想

一日中考えてなんとかわかったけど，本番で出たら絶対書けないと思いますね…．
DPは脳死でコードが書けないのでとてもつらいです．