SRM 666 Div2 Hard CollectingTokens

TopCoder Statistics - Problem Statement

• 復習。復習。
• 本番全く手が付けられなかったので、解説記事を見て書いた。
• 木の探索途中で部分 dp をやる。
• 与えられたグラフが非有向グラフなので par!=u の判定が必要なんだけど、dp 側でうっかりわすれてバグった。
• 葉ノードの判定で m==0 のケースが抜けて WA してしまった。
• こういうコードを簡単に書いて、一発で通せるようになりたい。

class CollectingTokens {
public:
int maxTokens(vector <int> A, vector <int> B, vector <int> tokens, int L) {
        int n = A.size()+1;
        vector<vector<int>> tree(n);
        REP(i, n-1)
        {
            tree[A[i]-1].push_back(B[i]-1);
            tree[B[i]-1].push_back(A[i]-1);
        }
        //
        // dp[v][l][b] := 頂点 v からたどるとき、残りの移動可能回数が l のときの最大取得トークン数
        //                ただし b==1 のときは v に戻ってくる
        //
        // 一部の子ノードには進まないという選択もある。
        //
        vector<vector<vector<int>>> dp(n,vector<vector<int>>(L+1,vector<int>(2,0)));
        function<void(int,int)> dfs = [&](int v, int par) {
            int m = tree[v].size();
            if (m == 0 || (m == 1 && tree[v][0] == par)) // 葉ノード(tree が 1 node だけのケースもある)
            {
                REP(b,2)
                REP(l, L+1)
                    dp[v][l][b] = tokens[v];
                return;
            }
            // 子の dp を更新
            for (auto u : tree[v])
                if (u != par)
                    dfs(u, v);
            // 自身の dp を更新
            //  * 一方通行の子は一つのみ。それはどれか？
            //  * どの子をたどるか？
            // dpdp[i][j][l] := 残りの移動可能回数が l のときi 番目までの子を見た時、
            //                j 番目で一方通行になる場合の最大取得トークン数
            vector<vector<vector<int>>> dpdp(m+1,vector<vector<int>>(m+1,vector<int>(L+1,0)));

            REP(i,m)   // m までの子のうちどれをたどるか？
            REP(j,m+1)   // どれが一方通行になるか？
            REP(k,L+1) // 残りの移動回数
            {
                int u = tree[v][i];
                dpdp[i+1][j][k] = max(dpdp[i+1][j][k], dpdp[i][j][k]); // i 番目を訪れないケース
                if (u == par) // 親にもどるものは飛ばす
                    continue;
                if (j < m) // 一方通行になるケース
                {
                    if (i == j) // j 番目で一方通行になる
                    {
                        for (int kk = 1; kk <= k; ++kk)
                            dpdp[i+1][j][k] = max(dpdp[i+1][j][k], dpdp[i][j][k-kk]+dp[u][kk-1][0]);
                    }
                    else
                    {
                        for (int kk = 2; kk <= k; ++kk)
                            dpdp[i+1][j][k] = max(dpdp[i+1][j][k], dpdp[i][j][k-kk]+dp[u][kk-2][1]);
                    }
                }
                else // 一方通行にならないケース
                {
                    for (int kk = 2; kk <= k; ++kk)
                        dpdp[i+1][m][k] = max(dpdp[i+1][m][k], dpdp[i][m][k-kk]+dp[u][kk-2][1]);
                }
            }
            REP(i,m+1)
            REP(k,L+1)
            {
                int b = (i==m)? 1 : 0;
                dp[v][k][b] = max(dp[v][k][b], dpdp[m][i][k] + tokens[v]);
            }
            return;
        };
        dfs(0, -1);
        return dp[0][L][0];
}
};