Typical DP Contest - P うなぎ

P: うなぎ - Typical DP Contest | AtCoder

  • dp の更新のためにさらに dp を使う問題。
const int MOD = (int)(1e+9)+7;
class p_unagi {
public:
    //
    // 根を一つ決めて有向木にして考える
    //
    // dp[x][i][j] := x を頂点とする部分木からパスを i 本選び、x での次数(xから出て行くパスの辺の本数)が
    //                j になっているようなパスの引き方
    //
    // x の子 y, z, w にたいして
    //
    //  1. y-x-z と x をまたぐパスを作成する
    //    dp[x][k+l+m-1][2] += dp[y][k][1]*dp[z][l][1]*dp[w][m][*]
    //  2 x-y パスを作成する
    //    dp[x][k+l+m][1] += dp[y][k][1]*dp[z][l][*]*dp[w][m][*]
    //  3. x を含まないパスを作成
    //    dp[x][k+l+m][0] += dp[y][k][*]*dp[z][l][*]*dp[w][m][*]
    //
    // のような更新となる。
    // このような組み合わせを計算するのに更に DP をする
    //
    vector<vector<vector<ll>>> dp_;
    vector<vector<int>> tree_;
    void dfs(int x, int par, int K) {
            vector<int> child;
            for (auto v : tree_[x])
            {
                if (v == par)
                    continue;
                child.push_back(v);
                dfs(v, x, K);
            }
            vector<vector<vector<ll>>> dp2;
            mdv::resize(dp2,child.size()+1,K+1,3,0);
            //
            // 通常の DP。i 番目までみたときにどのパスが選ばれているかではなくて、
            // 何本パスが選ばれているかを状態としてもっておけばよい。
            //
            // dp2[i][j][k] := i 番目の子まで見た時に j 個のパスがすでに選ばれていて
            //                 親につながるパスの辺の本数が k 本選ばれている状態の組み合わせ数
            //
            dp2[0][0][0] = 1;
            REP(i, child.size())
            REP(j, K+1)
            REP(k, 3)
            {
                // child[i] の結果を回る
                REP(l, K+1)
                REP(m, 3)
                {
                    ll way = dp2[i][j][k]*dp_[child[i]][l][m] % MOD;

                    // type 1.
                    // 親につなぐ(二本のパスをつないで1本にするので-1)
                    //   x           x
                    //  /           / \   :
                    // y   cur =>  y  cur
                    // |    |      |   |
                    // :    :      :   :
                    if (k == 1 && m == 1 && 0 <= j+l-1 && j+l-1 <= K)
                        (dp2[i+1][j+l-1][2] += way)%=MOD;

                    //   x           x
                    //  /           / \   :
                    // y   cur  => y  cur
                    // |           |
                    // :           :   y1
                    if (k == 1 && m == 0 && j+l <= K)
                        (dp2[i+1][j+l][k+1] += way)%=MOD;

                    // type 2.
                    // 親 x につなぐ(1本を親まで伸ばす)
                    //   x         x
                    //      =>      \    :
                    // y  cur    y  cur
                    //     |         |
                    //     :         :
                    if (k == 0 && m == 1 && j+l <= K)
                        (dp2[i+1][j+l][k+1] += way)%=MOD;
                    // x と y にだけパスを追加するケース(+1)
                    //
                    //   x             x
                    //                  \     :
                    // y   cur   =>  y   cur
                    // |             |
                    // :             :    y1
                    if (k == 0 && m == 0 && j+l+1 <= K)
                        (dp2[i+1][j+l+1][k+1] += way)%=MOD;

                    // type 3.
                    // 親 x につながないケース
                    if (j+l <= K)
                        (dp2[i+1][j+l][k] += way)%=MOD;
                }
            }
            // 結果を格納
            REP(i, K+1)
            REP(j, 3)
                dp_[x][i][j] = dp2[child.size()][i][j] % MOD;
            return;
    }
    void solve(void) {
            int N, K;
            cin>>N>>K;
            tree_.resize(N);
            REP(i, N-1)
            {
                int a, b;
                cin>>a>>b;
                --a;
                --b;
                tree_[a].push_back(b);
                tree_[b].push_back(a);
            }
            mdv::resize(dp_,N,K+1,3,0);

            // 大体 O(N*K^2*9) < 25000000 くらい
            // 根を 0 として dfs
            dfs(0, -1, K);
            ll res = 0;
            REP(i, 3)
                (res += dp_[0][K][i]) %= MOD;
            cout<<res<<endl;
    }
};