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

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

JOI 二次予選 2023 C - 塗りつぶし (AOJ 0749) (1D, 難易度 7)

ちょっと実装が重たい。

問題概要

 H \times W グリッドが与えられる。最初、どのマスにもなんらかの色(整数値)がついている。

マス  (i, j) から辺で接しているマスへの移動を繰り返し、マス  (i, j) と色が異なるマスに入ることなく移動できるマスの集まりを、マス  (i, j) の領域とよぶ。

今、いずれかのマスを 1 つ選んで、いずれかの色に塗りつぶす。このとき、そのマスの領域のマスもすべてその色に塗りつぶされる。

操作後において、そのマスの領域に含まれるマス数の最大値を求めよ。

制約

  •  1 \le H, W \le 500
  •  1 \le A_{i, j} \le 10^{9}

考えたこと

グリッドの各マスは「同じ色が上下左右に隣接しているマスを連結している」とみなしたときの、連結成分によって分類できる。このとき、次のようなグラフを考えよう。このグラフはたとえば Union-Find を用いて  O(HW \alpha(HW)) の計算量で構築できる。


  • 各連結成分を頂点とする
    • その頂点の重さは、連結成分に含まれるマス数
    • その頂点の色は、その連結成分に含まれるマスの色
  • 2 つの連結成分は、それぞれの連結成分に含まれるマスであって、互いに隣接しているものがあるとき、辺を張る

もとの問題は、こうして作られたグラフにおいて、頂点  v を 1 つ選び、その頂点の色を好きな色  c に塗り替えたとき、

  • その頂点  v の重さ
  • その頂点  v と隣接している頂点のうち、色が  c であるものの重さの総和

の和を最大化せよ、と解釈できる。

グラフの構築さえできてしまえば、各頂点について、隣接頂点を色ごとに分類する(たとえば map などを用いる)ことで解ける。

map を用いた場合、全体の計算量は  O(HW(\log H + \log W)) となる。

コード

#include <bits/stdc++.h>
using namespace std;
using pint = pair<int, int>;

// Union-Find
struct UnionFind {
    // core member
    vector<int> par, nex;

    // constructor
    UnionFind() { }
    UnionFind(int N) : par(N, -1), nex(N) {
        init(N);
    }
    void init(int N) {
        par.assign(N, -1);
        nex.resize(N);
        for (int i = 0; i < N; ++i) nex[i] = i;
    }
    
    // core methods
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool same(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
        par[x] += par[y];
        par[y] = x;
        swap(nex[x], nex[y]);
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

int main() {
    int H, W;
    cin >> H >> W;
    vector<vector<int>> A(H, vector<int>(W));
    for (int i = 0; i < H; i++) for (int j = 0; j < W; j++) cin >> A[i][j];

    // 連結成分を取り出していく
    UnionFind uf(H * W);
    for (int i = 0; i < H; i++) for (int j = 0; j < W; j++) {
        if (i+1 < H && A[i][j] == A[i+1][j]) uf.merge(i*W + j, (i+1)*W + j);
        if (j+1 < W && A[i][j] == A[i][j+1]) uf.merge(i*W + j, i*W + (j+1));
    }
    auto getroot = [&](int i, int j) -> int { return uf.root(i*W + j); };

    // 連結成分ごとの隣接関係を整理する
    set<pint> S;
    map<int, vector<pint>> G;     // G[v] := (隣接する連結成分の root, その色) を集めた隣接リスト
    for (int i = 0; i < H; i++) for (int j = 0; j < W; j++) {
        if (!G.count(getroot(i, j))) G[getroot(i, j)] = vector<pint>();
        if (i+1 < H && A[i][j] != A[i+1][j]) {
            if (!S.count(pint(getroot(i, j), getroot(i+1, j)))) {
                S.insert(pint(getroot(i, j), getroot(i+1, j)));
                S.insert(pint(getroot(i+1, j), getroot(i, j)));
                G[getroot(i, j)].emplace_back(getroot(i+1, j), A[i+1][j]);
                G[getroot(i+1, j)].emplace_back(getroot(i, j), A[i][j]);
            }
        }
        if (j+1 < W && A[i][j] != A[i][j+1]) {
            if (!S.count(pint(getroot(i, j), getroot(i, j+1)))) {
                S.insert(pint(getroot(i, j), getroot(i, j+1)));
                S.insert(pint(getroot(i, j+1), getroot(i, j)));
                G[getroot(i, j)].emplace_back(getroot(i, j+1), A[i][j+1]);
                G[getroot(i, j+1)].emplace_back(getroot(i, j), A[i][j]);
            }
        }
    }

    // 各連結成分について、隣接する連結成分を順に見ていく
    int res = 0;
    for (auto [node, adj] : G) {
        // 隣接する連結成分を色ごとに分類し、そのサイズを足していく
        map<int, int> ma;
        for (auto [v, color] : adj) ma[color] += uf.size(v);

        // そのサイズの最大値を求める
        int tmp = 0;
        for (auto [color, siz] : ma) tmp = max(tmp, siz);
        res = max(res, tmp + uf.size(node));
    }
    cout << res << endl;
}