mayoko’s diary

プロコンとかいろいろ。

SRM 666 div2 hard:CollectingTokens

やっぱ div2 hard も練習しないとダメですね(easy でごまかせたと思ったらこの問題で死んだ)。

解法

木に関するメモ化再帰的 dp をします。狙いとしては,

v 以下の部分木で移動回数が L に制限された時に最大でどれだけのトークンを得られるか?

というような dp をしたいです。つまり, dp[v][L] = (max 値) というような感じです。ただ, これだけだと上手く dfs に落としこむことが出来ないのでいろいろ工夫します。

構成としては, v の子について計算結果 dp[child][subL] を得てそれを利用して dp[v][L] を更新したい感じですが, 動き方として「v から child に移動して, v に返ってこない」というような動きも許されるので, フラグとして, 「v に返ってこないといけない」という変数を加えます。「v に返ってこないといけない」というフラグが立っているときは, v -> child -> (なにかする) -> child -> v -> (なにかする) という動きにしないといけない, というわけですね。

また, その他にも「親 v が子供 child の何番目まで移動するかどうかチェックしたか」という変数も必要です。よって, 最終的には

dp[v][e][L][m] = (頂点 v 以下の部分木で, v の e 番目までの子を調べ済みで, 残り移動回数は L 回で, 上で説明したようなフラグが m の時の最大トークン数)という dp を解くことになります。
dp の遷移はソースコードを見て欲しいです。

得た知見
  • やっぱり木の問題は部分木の部分問題に落としこむのが大事
    • 「どう落としこむか」を考えよう
  • 子について調べる時「どこまで調べたか」を管理する変数 e を使うと上手く行くことがある
const int MAXN = 55;
const int MAXL = 101;
vector<int> G[MAXN];
int dp[MAXN][MAXN][MAXL][2];
vector<int> tokens;

int dfs(int v, int e, int L, int m) {
    int& ret = dp[v][e][L][m];
    if (ret >= 0) return ret;
    ret = 0;
    if (e < (int)G[v].size()) {
        int u = G[v][e];
        for (int sL = 0; sL <= L; sL++) {
            // u に訪問したまま返ってこない
            if (sL + 1 <= L && m == 0) {
                int tmp = tokens[u] + dfs(u, 0, sL, m) + dfs(v, e+1, L-sL-1, 1);
                ret = max(ret, tmp);
            }
            // u に訪問し帰ってくる
            if (sL + 2 <= L) {
                int tmp = tokens[u] + dfs(u, 0, sL, 1) + dfs(v, e+1, L-sL-2, m);
                ret = max(ret, tmp);
            }
        }
        // u に訪問しない
        int tmp = dfs(v, e+1, L, m);
        ret = max(ret, tmp);
    }
    return ret;
}

class CollectingTokens {
public:
    int maxTokens(vector <int> A, vector <int> B, vector <int> tokens, int L) {
        ::tokens = tokens;
        int n = tokens.size();
        for (int i = 0; i < n; i++) G[i].clear();
        // make tree
        {
            set<int> st;
            queue<int> q;
            q.push(0);
            st.insert(0);
            while (!q.empty()) {
                int x = q.front(); q.pop();
                for (int i = 0; i < n-1; i++) {
                    int y = -1;
                    if (A[i]-1 == x) {
                        y = B[i]-1;
                    }
                    if (B[i]-1 == x) {
                        y = A[i]-1;
                    }
                    if (y != -1 && st.count(y)==0) {
                        G[x].push_back(y);
                        st.insert(y);
                        q.push(y);
                    }
                }
            }
        }
        // dfs
        memset(dp, -1, sizeof(dp));
        return tokens[0] + dfs(0, 0, L, 0);
    }
};