けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 097 F - Monochrome Cat (800 点)

とにかく重たい...

問題へのリンク

問題概要

 N 頂点のツリーが与えられる。各頂点には「白」か「黒」の色が塗られている。好きな頂点から開始して

  • 今いる頂点の色を flip する
  • 隣接する頂点を 1 つ選んで移動して、その頂点の色を flip する

といういずれかの操作を繰り返して、すべての頂点が黒色となるようにしたい。そのために必要な操作回数の最小値を求めよ。

制約

  •  1 \le N \le 10^{5}

解法 1: 全方位木 DP

まず、開始頂点  r を 1 つ固定した状態ならば、比較的簡明な木 DP で解くことができる。 r を根とした根付き木を考える。

  • dp1[ v ] := v を根とした部分木において、v からスタートして (スタート時に v は色反転するものとする)、v に戻って全部黒色にする最小回数
  • dp2[ v ] := v を根とした部分木において、v からスタートして (スタート時に v は色反転するものとする)、v に戻らなくてもよいが全部黒色にする最小回数

としよう。答えは、r の色を予め反転しておいた上での dp2[ r ] となる。

dp1 の遷移

遷移は以下の通り。

  • dp1[ v ] = sum(dp1[ c ], c: v の子) + 2c + parity1

parity1 部分は、v の子供の個数を n(v) として、「白」を 1、「黒」を 0 として

  • (v の初期色反転) + n(v) が偶数のとき parity1 = 0
  • (v の初期色反転) + n(v) が奇数のとき parity1 = 1

dp2 の遷移

遷移は以下の通り。

  • dp2[ v ] = sum(dp1[ c ], c: v の子) - max(dp1[ c ] - dp2[ c ], c: v の子) + 2c - 1 + parity2

parity2 は同様に、

  • (v の初期色反転) + n(v)-1 が偶数のとき parity2 = 0
  • (v の初期色反転) + n(v)-1 が奇数のとき parity2 = 1

全方位木 DP へ

以上の DP を、全頂点を根としてやりたい。そういうときは全方位木 DP にするとうまくいくイメージ。全方位木 DP は、

  • dp[ v ][ i ] := 頂点 v に接続している辺のうち i 番目の辺 (親含む) について、辺 (v, i) を除去することでできる 2 つの部分木のうち、i 側の木についての答え

という DP テーブルを作り上げる方法論である。元の木 DP ができていれば、全方位木 DP のフレームワークに乗っけると、ほぼ自動的に書くことができる。ここを参照。

罠に注意

ただし上記の DP は、「最初からすべての黒色」であるような部分木に対して適用するとおかしなことになってしまう。なぜなら、上記の DP はどんな場合であっても、葉まで行って「黒色」に塗ろうとするため...

よって、あらかじめ、グラフの葉から見て、後退解析して、黒色になっているところを削り落とす処理を行っておく。

そしてさらに、「すべての頂点が黒色の場合」がコーナーケースとなっていて、0 と答える必要がある。

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;

template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
using pll = pair<long long, long long>;
using Graph = vector<vector<int>>;
vector<int> color;
string str_color;
int N;
Graph origG, G;

// no-need subtree
vector<int> deg;
vector<bool> nonneed;

// dp, left, right (in 全方位木dp)
vector<vector<pll>> dp;
vector<vector<pll>> Left, Right;
vector<int> parent_ord; // G[v][parent_ord[v]] is v's parent
const pll CUMUL_UNITY = {0, 0};

pll merge(pll a, pll b) {
    return pll(a.first + b.first,
               a.first + b.first - max(a.first - a.second, b.first - b.second));
}

pll calc(int v, int jogai, int ini) {
    int s = G[v].size();
    int adjnum = s - 1;
    if (jogai == -1) ++adjnum;

    // leaf
    if (adjnum == 0) {
        if (ini % 2 == 1) return {1, 1};
        else return {0, 0};
    }

    // cumulative sum (not necessary to change)
    auto tmp = CUMUL_UNITY;
    if (adjnum == 1) {
        if (s == 1 || jogai == 1) tmp = dp[v][0];
        else tmp = dp[v][1];
    }
    else {
        if (jogai == -1) tmp = Left[v][s];
        else tmp = merge(Left[v][jogai], Right[v][s-jogai-1]);
    }
    
    // merge
    auto res = tmp;
    res.first += adjnum * 2;
    res.second += adjnum * 2 - 1;
    if ( (ini + adjnum) % 2 == 1 ) ++res.first;
    if ( (ini + adjnum + 1) % 2 == 1 ) ++res.second;
    
    return res;
}

void build(int v) {
    int s = G[v].size();
    Left[v].assign(s+1, CUMUL_UNITY);
    Right[v].assign(s+1, CUMUL_UNITY);
    for (int i = 0; i < s; ++i) {
        Left[v][i+1] = merge(Left[v][i], dp[v][i]);
        Right[v][i+1] = merge(Right[v][i], dp[v][s-i-1]);
    }
}
    
pll rec(int v, int p = -1) {
    int s = G[v].size();
    dp[v].assign(s, CUMUL_UNITY);
    for (int i = 0; i < s; ++i) {
        int to = G[v][i];
        if (to == p) {
            parent_ord[v] = i;
            continue;
        }
        dp[v][i] = rec(to, v);
    }
    build(v);
    return calc(v, parent_ord[v], 1 - color[v]);
}

void rerec(int v, pll pval, int p = -1) {
    int s = G[v].size();
    if (parent_ord[v] != -1) dp[v][parent_ord[v]] = pval;
    build(v);
    for (int i = 0; i < s; ++i) {
        int to = G[v][i];
        if (to == p) continue;
        auto nval = calc(v, i, 1 - color[v]);        
        rerec(to, nval, v);
    }
}

void pre() {
    nonneed.assign(N, false);
    queue<int> que;
    for (int v = 0; v < N; ++v) if (color[v] == 0 && deg[v] == 1) que.push(v);
    while (!que.empty()) {
        int v = que.front(); que.pop();
        nonneed[v] = true;
        for (auto to : origG[v]) {
            if (nonneed[to]) continue;
            --deg[to];
            if (color[to] == 0 && deg[to] == 1) que.push(to);
        }
    }
    for (int v = 0; v < N; ++v) {
        if (nonneed[v]) continue;
        for (auto to : origG[v]) {
            if (nonneed[to]) continue;
            G[v].push_back(to);
        }
    }
}

long long solve() {
    // cut down black leaves
    pre();

    // DP
    dp.assign(N, vector<pll>());
    Left.assign(N, vector<pll>());
    Right.assign(N, vector<pll>());
    parent_ord.assign(N, -1);

    // 後退解析で削られなかった頂点を探す
    int r = 0;
    for (; r < N; ++r) if (!nonneed[r]) break;
    if (r == N) return 0;
    
    rec(r);
    rerec(r, {0, 0});

    // search
    long long res = INF;
    for (int v = 0; v < N; ++v) {
        if (nonneed[v]) continue;
        auto tmp = calc(v, -1, color[v]);
        chmin(res, tmp.second);
    }
    return res;
}

int main() {
    cin >> N;
    origG.assign(N, vector<int>());
    G.assign(N, vector<int>());
    deg.assign(N, 0);
    for (int i = 0; i < N-1; ++i) {
        int u, v; cin >> u >> v; --u, --v;
        origG[u].push_back(v), origG[v].push_back(u);
        ++deg[u], ++deg[v];
    }
    cin >> str_color;
    color.assign(N, 0);
    for (int i = 0; i < N; ++i) if (str_color[i] == 'W') color[i] = 1;
    cout << solve() << endl;
}

解法 2:直径を求めるのに類似した方法

解法 1 の方法は、全方位木 DP でありがちな「根から各頂点への距離の最大値が最大となる根を求める」というのによく似ている。これは直径を求めることで実現できるのだが、今回の問題もそれに類似した方法 (DFS 2 回) で求めることもできる。公式解説参照。

https://img.atcoder.jp/arc097/editorial.pdf