Codeforces Round #322 (Div.2) F. Zublicanes and Mumocrates

問題文

http://codeforces.com/contest/581/problem/F

解法

（知っていれば）二乗の木DPっぽいな～となり、実際そうである。
ひとまず普通に木DPをする。
dp[v][cnt][c] := v の部分木だけ考えて、その中の葉を cnt だけ黒く塗り、v の色が c であるときの最小値
とする。
遷移は
dp[v][cnt][c] := min(dp[v][cnt][c], dp[v][cnt - ch_cnt][c] + dp[to][ch_cnt][ch_c] + (c != ch_c)) (forall ch)
みたいなことをする。ただし、v が葉でない場合は、ひとまず dp[v] の初期値を子から1つ選んで初期化しておく必要がある。
以降は最初に選んだ子は無視して遷移する。
遷移するときに、今処理してる子による結果が影響しないようにしておくこと。

さて、問題は計算量だが、cnt と ch_cnt によるループで各ノードで O(n^2) かかるので、全体として O(n^3) っぽいが、実際はそれほど計算が発生しない。

かなり大雑把に捉えるなら、例えば2つの子を持つ部分木を考えて、子のサイズを l, r としたとき、T(l + r) = T(l) + T(r) + l * r のような計算量の漸化式を考えると良い。
(l + r)^2 = l^2 + r^2 + l * r なので、T(n) = n^2 とみなせることがなんとなくわかるはず。

あとはMLEに注意する。あと、ループはきっちりサイズちょうどの上限で回さないと計算量が変わるので丁寧にやること。

ソースコード

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

constexpr int inf = 1e9;

int main() {
    int n; cin >> n;
    vector<vector<int>> g(n);
    for(int i = 0; i < n - 1; ++i) {
        int a, b; cin >> a >> b;
        g[a - 1].push_back(b - 1);
        g[b - 1].push_back(a - 1);
    }
    if(n == 2) {
        cout << 1 << endl;
        return 0;
    }
    int root = -1;
    for(int i = 0; i < n; ++i) {
        if(g[i].size() != 1u) root = i;
    }

    vector<vector<vector<int>>> dp(n, vector<vector<int>>(2, vector<int>(2, inf)));
    vector<int> sz(n);
    function<void(int, int)> dfs = [&] (int v, int p) {
        if(g[v].size() == 1u) {
            sz[v] = 1;
            dp[v][0][0] = dp[v][1][1] = 0;
            return;
        }
        int fst_ch = -1;
        for(auto to : g[v]) {
            if(to == p) continue;
            if(fst_ch == -1) fst_ch = to;
            dfs(to, v);
        }
        {
            int sum = 0;
            for(auto to : g[v]) {
                if(to == p) continue;
                sum += sz[to];
            }
            dp[v].assign(sum + 1, vector<int>(2, inf));
        }
        for(auto to : g[v]) {
            if(to == p) continue;
            sz[v] += sz[to]; // !! add here
            if(to == fst_ch) {
                for(int cnt = 0; cnt <= sz[v]; ++cnt) {
                    for(int c = 0; c < 2; ++c) {
                        for(int cc = 0; cc < 2; ++cc) {
                            if(dp[to][cnt][cc] == inf) continue;
                            dp[v][cnt][c] = min(dp[v][cnt][c], dp[to][cnt][cc] + (c != cc));
                        }
                    }
                }
            } else {
                for(int cnt = sz[v]; cnt >= 0; --cnt) {
                    for(int c = 0; c < 2; ++c) {
                        int mini = inf;
                        for(int ch_cnt = 0; ch_cnt <= min(sz[to], cnt); ++ch_cnt) {
                            for(int ch_c = 0; ch_c < 2; ++ch_c) {
                                mini = min(mini, dp[v][cnt - ch_cnt][c] + dp[to][ch_cnt][ch_c] + (c != ch_c));
                            }
                        }
                        dp[v][cnt][c] = mini;
                    }
                }
            }
        }
    };
    dfs(root, -1);

    cout << min(dp[root][sz[root] / 2][0], dp[root][sz[root] / 2][1]) << endl;
}

感想

二乗の木DP、集中してコーディングしないとすぐバグるからつらい。