Union-Find を使いこなす！！！

問題概要 (意訳)

$K$ 個の 0-1 変数 $x_{0}, x_{1}, \dots, x_{K-1}$ が与えられていて、最初はそれらの値について特に制約はない。いま、 $N$ 個の制約が順に与えられる。各制約はそれぞれ

1 つの変数 $x_{i}$ と 0 か 1 の値 w を指定して、 $x_{i} = w$ とする
2 つの変数 $x_{i}, x_{j}$ と 0 か 1 の値 w を指定して、 $x_{i} + x_{j} ≡ w \bmod 2$ とする

という形式をしている。 $k = 1, 2, \dots, N$ について、 $k$ 番目の制約を追加した時点でとりうる $x_{0}, \dots, x_{K-1}$ の値の組のうち、1 の個数の最小値を求めよ。

制約

$1 \le K, N \le 10^{5}$

解法分類

ざっくり以下についてどう扱うかによって、様々な解法が考えられそう

2 つの変数の偶奇を等しくするか、異なるようにするかの制約
1 つの変数の値を fix する
クエリ先読みするかどうか

1 個目については「重み付き Union-Find を使う」or「頂点倍加 Union-Find を使う」といった感じ。

2 個目については、重み付き Union-Find であれば、まともに取り扱うこともできる。頂点倍加 Union-Find の場合は「スラック頂点」を用意してあげると明快に記述できる。

3 個目については、クエリ先読みをすることで、各変数間の相対的な関係が最終的にどうなるのかが予めわかれば、色々楽できるという話に。

解法(1)：値 fix 付きの重み付き Union-Find で、前から解く

コンテスト本番では、重み付き Union-Find を使って気合いで前から解いた。コンテスト後にクエリ先読みして、先に各変数の相対的な関係性を把握してからクエリ処理した方が楽になることを学んだ。

なにはともあれ、まずは前から力技で解いてみる。通常の重み付き Union-Find は、各要素 v に重み x[ v ] をもたせて、

x[ i ] - x[ j ] = d となるように i と j をつなぐ

といったことができるようにしたもの。重みを偶奇にすれば OK。これで制約のうちの 2 番目を扱える。今回はさらに

x[ i ] の値を 0 または 1 に指定する

というのを扱えるようにする必要がある。そこで重み付き Union-Find に、各連結成分の根 r に対して

val[ r ] := 0 or 1 or -1 (-1 の場合はそのグループは差分のみわかっている状態で未確定)

という値を持たせることにした。これは次のようにすれば OK。

set 時

変数 x を値 w にセットする。このとき、まずは x の根 r を求めながら、x と r との diff を計算する。

そしてこの diff で w を補正してあげる気持ちで、val[ r ] = w ^ diff とすれば OK

merge 時

変数 x と変数 y を、通常の重み付き Union-Find と同様に、距離 w を補正しながら根に持ってくる。マージテクで x と y を適宜入れ替えて、x を根にする状態にする (y の親を x にする)

このとき、もし val[ y ] != -1 であったならば、val[ x ] = val[ y ] ^ w とする。

さらに「1 の個数の最大値」

さらに一工夫必要になる。Union-Find の各連結成分の根 r に対して、

onum[ r ] := 仮に x[ r ] = 0 としたときに、その連結成分内に 1 が何個あるか

という情報をもたせて、適宜更新するようにする。これを持っておけば、その連結成分について

val[ r ] = -1 のとき、min(onum[ r ], -par[ r ] - onum[ r ])　(-par[ r ] は連結成分のサイズ)
val[ r ] = 0 のとき、onum[ r ]
val[ r ] = 1 のとき、-par[ r ] - onum[ r ]

によって「 1 の個数として考えられる最小値」が得られる。毎回のクエリ更新では、更新が発生しうる連結成分についてのみ、更新前後の「1 の個数の最大値」を求めてあげて、その差分を求めれば OK。

以上で計算量は $O((N + K) \alpha(K) )$ となる。

#include <iostream>
#include <vector>
#include <map>
using namespace std;

template<class Abel> struct UnionFind {
    const Abel UNITY_SUM = 0;      // to be set
    vector<int> par;
    vector<Abel> diff_weight;
    vector<Abel> val;
    vector<int> onum; // 根が 0 のときの 1 の個数
    
    UnionFind(int n) : par(n, -1), diff_weight(n, UNITY_SUM)
                     , val(n, -1), onum(n, 0) {}
    int root(int x) {
        if (par[x] < 0) return x;
        else {
            int r = root(par[x]);
            diff_weight[x] ^= diff_weight[par[x]];
            return par[x] = r;
        }
    }
    
    Abel calc_weight(int x) {
        int rx = root(x);
        return diff_weight[x];
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }

    void set(int x, Abel w) {
        auto rx = root(x);
        auto dw = diff_weight[x];
        val[rx] = w ^ dw;
    }
    
    bool merge(int x, int y, Abel w = 0) {
        w ^= calc_weight(x); w ^= calc_weight(y);
        x = root(x), y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        
        if (w == 0) onum[x] += onum[y];
        else onum[x] += -par[y] - onum[y];

        if (val[y] != -1) val[x] = val[y] ^ w;
        par[x] += par[y];
        par[y] = x;
        diff_weight[y] = w;
        return true;
    }
    
    Abel diff(int x, int y) {
        return calc_weight(y) ^ calc_weight(x);
    }
    
    int size(int x) {
        return -par[root(x)];
    }

    int get_onum(int x) {
        x = root(x);
        if (val[x] == -1) return min(onum[x], -par[x] - onum[x]);
        else if (val[x] == 0) return onum[x];
        else return -par[x] - onum[x];
    }
};

/* debug */
template<class Abel> ostream& operator << (ostream& s, UnionFind<Abel> uf) {
    map<int, vector<int> > res;
    cout << endl;
    for (int i = 0; i < uf.par.size(); ++i) {
        int r = uf.root(i);
        res[r].push_back(i);
        if (r == i) {
             cout << "root " << r << ": "
                  << uf.val[r] << ", " << uf.onum[r] << endl;
        }
    }
    int iter = 0;
    for (auto it : res) {
        s << (iter++) << "-th group: ";
        for (int j = 0; j < (int)it.second.size(); ++j) {
            s << it.second[j] << "(" << uf.calc_weight(it.second[j]) << "), ";
        }
        s << endl;
    }
    return s << endl;
}


// 入力
int N, K;
string S;
vector<vector<int>> cls;

void solve() {
    long long res = 0;
    UnionFind<int> uf(K);
    
    for (int i = 0; i < N; ++i) {
        int w = 1 - (int)(S[i] - '0');
        int pre_one = -1;
        int cur_one = -1;
        
        if (cls[i].size() == 1) {
            int x = cls[i][0];
            pre_one = uf.get_onum(x);
            uf.set(x, w);
            cur_one = uf.get_onum(x);
            res += (cur_one - pre_one);
        }
        else if (cls[i].size() == 2) {
            int x = cls[i][0], y = cls[i][1];
            if (!uf.issame(x, y)) {
                pre_one = uf.get_onum(x) + uf.get_onum(y);
                uf.merge(x, y, w);
                cur_one = uf.get_onum(x);
                res += (cur_one - pre_one);
            }
        }
        cout << res << endl;
    }
}

int main() {
    cin >> N >> K >> S;
    cls.assign(N, vector<int>());
    for (int k = 0; k < K; ++k) {
        int c; cin >> c;
        for (int i = 0; i < c; ++i) {
            int a; cin >> a; --a;
            cls[a].push_back(k);
        }
    }
    solve();
}

解法(2)：スラック頂点を用意して、頂点倍加 Union-Find で前から解く

今回のように

変数 x と y を同じ色で塗る
変数 x と y を異なる色で塗る

というような制約は、二部グラフ判定に近い香りがある。よって、二部グラフ判定の典型技として知られる「頂点倍加 Union-Find」で取り扱うことができる。先に頂点倍加 Union-Find について復習してみる。

各頂点 v に対して、Union-Find 上で対応する頂点として

(v, 0)： x[ v ] を 0 にすることを表す頂点
(v, 1)： x[ v ] を 1 にすることを表す頂点

を用意する。頂点数は 2 倍になる。そして頂点 u と頂点 v を異なる値に塗ることを

(u, 0) と (v, 1) とをつなぎ、(u, 1) と (v, 0) とをつなぐ

という風にする。これは x[ u ] = 0 とするならば x[ v ] = 1 とすることを表し、x[ v ] = 1 とするならば x[ u ] = 0 とすることを表す。ここまでは二部グラフ判定でおなじみだけど、今回は「頂点 u と頂点 v を同じ値に塗る」とうい制約もある。これも同様に

(u, 0) と (v, 0) とをつなぎ、(u, 1) と (v, 1) とをつなぐ

とすればよい。さらに問題となるのは、「変数 x を 0 または 1 にする」という制約だ。でもこれに対しては、「そっちには固定するよ、という変数 s」と「そっちには固定しないよ、という変数 t」を用意して、

変数 v を 0 にするなら、s と (v, 0) とをつないで、t と (v, 1) とをつなぐ
変数 v を 1 にするなら、s と (v, 1) とをつないで、s と (v, 0) とをつなぐ
変数 s の初期重みは 0、変数 t の初期重みは INF

という風にすれば OK。また、毎回の更新時において「Union-Find 上の各連結成分における 1 の個数」は常に更新し続ける。

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

const int INF = 1<<29;
struct UnionFind {
    int s, t; // 動かざる頂点
    vector<int> par;
    vector<int> onum; // 連結成分内に 1 が何個あるか

    // v: 0 側, v + n: 1 側
    UnionFind(int n) : s(n), t(n * 2 + 1),
                       par(n * 2 + 2, -1), onum(n * 2 + 2, 0) {
        for (int i = n + 1; i < n * 2 + 1; ++i) onum[i] = 1;
        onum[t] = INF;
    }
    
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        
        // x を根に
        onum[x] += onum[y];
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }

    int get_onum(int x) {
        return onum[root(x)];
    }
};

// 入力
int N, K;
string S;
vector<vector<int>> cls;

void solve() {
    long long res = 0;
    UnionFind uf(K);

    // 変数 x を 0 にする場合と、1 にする場合の小さい方
    auto get = [&](int x) {
        x %= K+1;
        return min(uf.get_onum(x), uf.get_onum(x+K+1));
    };
    
    // 変数 x と y を結ぶ
    auto merge = [&](int x, int y) {
        if (uf.issame(x, y)) return;
        int before_cost = get(x) + get(y);
        uf.merge(x, y);
        if (x < K+1) x += K+1; else x -= K+1;
        if (y < K+1) y += K+1; else y -= K+1;
        uf.merge(x, y);
        int after_cost = get(x);
        res += after_cost - before_cost;
    };  
    
    for (int i = 0; i < N; ++i) {
        int w = 1 - (int)(S[i] - '0');
        int pre_one = -1;
        int cur_one = -1;
        
        if (cls[i].size() == 1) {
            int x = cls[i][0];
            if (w == 0) merge(x, uf.s);
            else merge(x+K+1, uf.s);
        }
        else if (cls[i].size() == 2) {
            int x = cls[i][0], y = cls[i][1];
            if (w == 0) merge(x, y);
            else merge(x, y+K+1);
        }
        cout << res << endl;
    }
}

int main() {
    cin >> N >> K >> S;
    cls.assign(N, vector<int>());
    for (int k = 0; k < K; ++k) {
        int c; cin >> c;
        for (int i = 0; i < c; ++i) {
            int a; cin >> a; --a;
            cls[a].push_back(k);
        }
    }
    solve();
}

解法(3)：クエリ先読みして、各変数の相対的関係を先に求めてしまう

さらにさらに、今回クエリを先読みすることで、各変数の相対的な関係性を先に知ることができる。それを利用すると、Union-Find で頂点倍加する必要すらなくなる。

「変数 v の値を 0 または 1 にする」というクエリに対して、スラック変数 s を用意して (v, s) をつなぐのは解法 (2) と同様。

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

const int INF = 1<<29;
struct UnionFind {
    int s; // 動かざる頂点
    vector<int> par;
    vector<int> onum; // 連結成分内に 1 が何個あるか
           
    UnionFind(int n, const vector<int> &v) : s(n), par(n+1, -1), onum(v) {}
           
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        
        // x を根に
        onum[x] += onum[y];
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }

    int get_onum(int x) {
        return onum[root(x)];
    }
};

// 入力
int N, K;
string S;
vector<vector<int>> cls;

using Edge = pair<int,int>;
using Graph = vector<vector<Edge>>;

void solve() {
    long long res = 0;
    vector<int> color(K+1, -1);
    Graph G(K+1);
    for (int i = 0; i < cls.size(); ++i) {
        int w = 0;
        if (S[i] == '0') w = 1;
        if (cls[i].size() == 1) {
            int x = cls[i][0];
            G[x].emplace_back(K, w), G[K].emplace_back(x, w);
        }
        else if (cls[i].size() == 2) {
            int x = cls[i][0], y = cls[i][1];
            G[x].emplace_back(y, w), G[y].emplace_back(x, w);
        }
    }
    auto dfs = [&](auto &&dfs, int v, int c = 0) -> void {
        color[v] = c;
        for (auto e : G[v]) {
            if (color[e.first] != -1) continue;
            dfs(dfs, e.first, c^e.second);
        }
    };
    // color[K] は優先的に 0 に
    for (int i = K; i >= 0; --i) if (color[i] == -1) dfs(dfs, i, 0);

    // 先読みした color を Union-Find に渡す
    UnionFind uf(K, color);

    // x が固定されていたら onum、固定されていなかったら min(onum, size-onum)
    auto get = [&](int x) {
        int num = uf.get_onum(x);
        if (uf.issame(x, K)) return num;
        else return min(num, uf.size(x) - num);
    };

    // 統合
    auto merge = [&](int x, int y) {
        if (uf.issame(x, y)) return;
        int before_cost = get(x) + get(y);
        uf.merge(x, y);
        int after_cost = get(x);
        res += after_cost - before_cost;
    };

    // 処理
    for (int i = 0; i < N; ++i) {
        if (cls[i].size() == 1) {
            int x = cls[i][0];
            merge(x, K);
        }
        else if (cls[i].size() == 2) {
            int x = cls[i][0], y = cls[i][1];
            merge(x, y);
        }
        cout << res << endl;
    }
}

int main() {
    cin >> N >> K >> S;
    cls.assign(N, vector<int>());
    for (int k = 0; k < K; ++k) {
        int c; cin >> c;
        for (int i = 0; i < c; ++i) {
            int a; cin >> a; --a;
            cls[a].push_back(k);
        }
    }
    solve();
}

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

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

Codeforces Round #616 (Div. 1) C. Prefix Enlightenment (R2400)

問題概要 (意訳)

制約

解法分類

解法(1)：値 fix 付きの重み付き Union-Find で、前から解く

set 時

merge 時

さらに「1 の個数の最大値」

解法(2)：スラック頂点を用意して、頂点倍加 Union-Find で前から解く

解法(3)：クエリ先読みして、各変数の相対的関係を先に求めてしまう

問題概要 (意訳)

制約

解法分類

解法(1)： 値 fix 付きの重み付き Union-Find で、前から解く

set 時

merge 時

さらに「1 の個数の最大値」

解法(2)： スラック頂点を用意して、頂点倍加 Union-Find で前から解く

解法(3)： クエリ先読みして、各変数の相対的関係を先に求めてしまう

解法(1)：値 fix 付きの重み付き Union-Find で、前から解く

解法(2)：スラック頂点を用意して、頂点倍加 Union-Find で前から解く

解法(3)：クエリ先読みして、各変数の相対的関係を先に求めてしまう