問題概要

頂点数 $N$ の木が与えられる。各頂点には 0, 1, 2 のいずれかの色 (0: 灰色, 1: 白色, 2: 黒色) が割り振られている。以下の操作を最小回数行うことで、木を消滅させたい。最小回数を求めよ。

[操作]：残っている頂点の部分集合 $S$ を選んで、それらとそれらに接続する辺をすべて削除する。ただし「 $S$ のどの 2 点も連結」「 $S$ の中に色 1 の頂点と色 2 の頂点を同時には含まない」という条件を満たさなければならない。

(というテストケースが $T$ 個与えられる)

制約

$1 \le N の総和 \le 2 \times 10^{5}$

考えたこと

最初はてっきり

最初に色 1 の頂点をすべて除去して (1 回)、残った連結成分それぞれを順に削除 (連結成分の個数回)
最初に色 2 の頂点をすべて除去して (1 回)、残った連結成分それぞれを順に削除 (連結成分の個数回)

のうちの小さい方が答えだと思ってしまった。しかし割とすぐに反例が見つかる。下図の場合、最初にすべての黒を除去しても、白を除去しても、大量の連結成分が生えてしまう。

この図の場合、正しくは 3 回でできる。

黒葉を除去する
白葉を除去する (この段階で白がすべて消える)
残りの黒を除去する

この例を見ると「同色の葉をすべて除去する」を繰り返すのが良さそうという直感が生える。より正確には次のようにする。

最初に除去する色を決め打ちする (色 1 か色 2)
色 1 と色 2 について交互に以下を繰り返す (木が消滅するまで)
- 該当する色 (色 0 も含む) の葉が存在する限り、それを削除することを繰り返す

実装上は queue を使った後退解析とかでできる。

葉からの除去でいいことの証明

木に関する問題を解く定石として、まずはパスについて考えるのがよさそう。パスについて考えると、とりあえず

白頂点が連続している箇所は縮約してよい
黒頂点が連続している箇所は縮約してよい
灰頂点は除外して縮約しても一緒

ということがわかる。よって、白黒が連続しているような状態に帰着される。これは頂点数を $N$ としたとき、どんなに上手にやっても $(N+1)/2$ 回未満にはできない (そして $(N+1)/2$ 回でできる) ことが帰納法から言える。 $(N+1)/2$ 回というのは、まさに「同色の葉から取り除く作業を色を交互に繰り返す」という作業を実施した場合の回数に一致している。

よって、パスグラフの場合には「同色の葉から取り除く」という操作が最適になることが言えた。一般のグラフに対しても「直径」を考えることでまったく同様に示される。

コード

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

const int INF = 1<<29;
using Graph = vector<vector<int>>;
int solve(int N, Graph G, vector<int> C, vector<int> deg, int col) {
    if (N == 1) return 1;
    int res = 0;
    queue<int> que, nque;
    for (int v = 0; v < N; ++v) {
        if (deg[v] == 1) {
            if (C[v] == col || C[v] == 0) que.push(v);
            else nque.push(v);
        }
    }
    if (que.empty()) return INF;

    vector<bool> seen(N, false);
    while (!que.empty()) {
        ++res;
        int ncol = 3 - col;
        while (!que.empty()) {
            int v = que.front(); que.pop();
            seen[v] = true;
            for (auto nv : G[v]) {
                if (seen[nv]) continue;
                --deg[nv];
                if (deg[nv] == 1) {
                    if (C[nv] == col || C[nv] == 0) que.push(nv);
                    else nque.push(nv);
                }
            }
        }
        swap(que, nque), swap(col, ncol);
    }
    return res;
}

int main() {
    cin.tie(0); 
    ios::sync_with_stdio(false);

    int T; cin >> T;
    while (T--) {
        int N;
        cin >> N;
        Graph G(N);
        vector<int> C(N), deg(N, 0);
        for (int i = 0; i < N; ++i) cin >> C[i];
        for (int i = 0; i < N-1; ++i) {
            int u, v; cin >> u >> v; --u, --v;
            G[u].push_back(v), G[v].push_back(u);
            ++deg[u], ++deg[v];
        }
        cout << min(solve(N, G, C, deg, 1), solve(N, G, C, deg, 2)) << endl;
    }
}