yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

E - Christmas Color Grid 1 / ユニークビジョンプログラミングコンテスト2023 クリスマス (AtCoder Beginner Contest 334)

問題

 h w 列のグリッドがあり、それぞれのマスは赤色または緑色に塗られている。

グリッドの上から  i 番目、左から  j 番目のマスを  (i, \, j) と表記する。 マス  (i, \, j) の色は文字  {s}_{i, j} で表され、  {s}_{i, j} = . のときは赤色、  {s}_{i, j} = # のときは緑色である。

グリッドにおいて、緑色に塗られたマスを頂点集合、隣り合った2つの緑色のマスを結ぶ辺全体を辺集合としたグラフにおける連結成分の個数を「緑の連結成分数」と呼ぶ。 ただし、2つのマス  (x, \, y) (x', \, y') が隣り合うとは、  |x - x'| + |y - y'| = 1 であることを指す。

赤色に塗られたマスをランダムに1つ選び、緑色に塗り替えるとき、塗り替え後のグリッド内の緑の連結成分数の期待値を  \mathrm{mod} \, 998244353 で出力せよ。 *1

入力

まず最初の1行目に整数  h, w が与えられる。

次の  h 行に渡り、  i 行目には文字列  {s}_{i, 1} {s}_{i, 2} \cdots {s}_{i, w} が順に与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  •  1 \leq h, w \leq 1000
  •  {s}_{i, j} = . または #
  •  {s}_{i, j} = . である  (i, \, j) が少なくとも1つは存在する。

出力

答えを1行で出力すること。

解法

マス  (i, \, j) が赤色に塗られていたとして、これを緑色に塗り替えるときに、グリッド全体の緑の連結成分数はどのようになるのかをまずは考えます。

たとえば、以下の図のように中央の赤いマスを塗り替えるときについて考えます。

このうち、元々右のマスと下のマスが緑の連結成分であった場合は、中央を緑に塗り替えても連結成分の数は変わりません。

一方で、右のマスと下のマスが緑の連結成分でない場合には、中央を緑に塗り替えたときに、これらは1つの連結成分となります。 すなわち、2つの連結成分だったものがなくなり、新たに大きな1つの連結成分になる、ということが言えます。

以上から、あるマスを緑色に塗り替えたときには、その上下左右に存在する連結成分がすべて1つの連結成分になると言えます。 従って、塗り替え後のグリッド全体の緑の連結成分数は「(塗り替え前のグリッドに存在している緑の連結成分数) -(マス  (i, j) に隣接している緑の連結成分数) + 1 」によって与えられます。

このことから、すべての緑色のマスについて、どの連結成分に所属しているかをまずは調べます。 同じ連結成分かどうかについては、

  1. 未探索の  {s}_{i, j} = # である  (i, \, j) を取り上げる
  2. その  (i, \, j) の上下左右のもののうち、緑色のものを次の探索対象とする
  3. 次の調査対象について、手順2を行い、探索を繰り返す
  4. 手順1で取り上げた  (i, \, j) を元にして探索できるものがなくなったら終了する

という手順によって探ることができます。

ただし、これを赤色のマスが出現するたびに行うと、毎回の計算が膨大になってしまいます。 そのため、予め緑色のマスに対してナンバリングを行うなどして、「同じ番号ならば同じ連結成分」ということを保存した状態にしておくと、同じ連結成分かどうかの判定は  O(1) で行うことができるようになります。

さて、ある1つの赤色のマスを緑色に塗り替えたときの、グリッド全体の連結成分数はこれで求めることができました。 ただ、この問題で求めるものは赤色のマスをランダムに選び、緑色に塗り替えたときの連結成分数の期待値です。 これは、 \begin{align} \frac{\text{(「ある赤色のマスを塗り替えた後の全体の連結成分数」の総数)}}{\text{(赤色のマスの総数)}} \end{align} によって求めることができます。 従って、それぞれの赤色のマスについて、塗り替え後の全体の連結成分数を求め、その総和を算出することによって、期待値の計算も行えると言えます。

以上を実装することにより、この問題の解答が得られます。 計算量としては、緑色のマスの連結成分について探索する部分で  O(hw) が、また各赤色のマスについて連結成分数を算出する部分で  O(hw) がかかり、これは実行時間制限に十分に間に合うと考えられます。

ソースコード

まず、マス  (i, \, j) が緑色である場合に、その上下左右のマスについても調べてナンバリングする関数 dfs() を以下のように実装しました。

char s[1005][1005]; // 文字の2次元配列で扱うことにより、 1-index で扱えるようにした
int group[1005][1005]; // 連結成分のナンバリングを保存する2次元配列。 group[i][j] = 0 ならば未探索
int group_count; // 連結成分の個数を表す整数。緑色で未探索のものが出現したら +1 される
const int di[4] = {-1, 0, 1, 0}, dj[4] = {0, -1, 0, 1}; // 上下左右の移動のために使用する配列

void dfs(int i, int j) {
  group[i][j] = group_count; // (i, j) のナンバリングを設定する

  // di, dj を用いて上下左右の4方向について探索する
  for (int k = 0; k < 4; k++) {
    if (group[i + di[k]][j + dj[k]] == 0 && s[i + di[k]][j + dj[k]] == '#') {
      dfs(i + di[k], j + dj[k]); // もし上下左右で未探索のマスが緑色ならば、そのマスにも dfs を行う
    }
  }
  return;
}

この dfs() によって、すべての緑色のマスに対して同じ連結成分かどうかが判定できる状態になったときに、赤色のマス  (i, \, j) を塗り替えた後の全体の連結成分数を求める関数 search() を以下のように実装しました。

int group[1005][1005]; // 連結成分のナンバリングを保存する2次元配列
int group_count; // 連結成分の個数を表す整数
const int di[4] = {-1, 0, 1, 0}, dj[4] = {0, -1, 0, 1}; // 上下左右の移動のために使用する配列

int search(int i, int j) {
  set<int> tmp; // 上下左右の緑色のマスのナンバリングを保存する set (set にしたのは、重複を扱うため)

  // 上下左右の4方向について、緑色のマスの連結成分の番号を調べる
  for (int k = 0; k < 4; k++) {
    if (group[i + di[k]][j + dj[k]] > 0) { // group の値が 0 より大きいならば、緑色のマス
      tmp.insert(group[i + di[k]][j + dj[k]]);
    }
  }

  return (group_count + 1) - tmp.size(); // (全体の連結成分数 + 1 - 上下左右の連結成分数)を返す
}

また、 以前の記事 でも触れた通り、既約分数  \displaystyle \frac{a}{b} \mathrm{mod} \, p は \begin{align} \frac{a}{b} \equiv {a} {b}^{p - 2} \quad ( \mathrm{mod} \, p) \end{align} によって求められます。 (ただし、  b p は互いに素であるとします。)

すなわち、この問題ではある整数の  (998244353 - 2) 乗を求めなければなりません。 従って、以下のような累乗を求める時間を短縮する my_pow() 関数を作成しました。

const ll MOD2 = 998244353;

ll my_pow(ll a, ll n) { // a の n 乗を求める
  if (n == 0) {
    return 1; // n = 0 ならば 1 を返す
  }

  if (n % 2 == 0) {
    return my_pow(a * a % MOD2, n / 2); // n が 2 で割り切れるならば、 (a ^ 2) ^ (n / 2) を計算する
  } else {
    return (a * my_pow(a, n - 1)) % MOD2; // n が 2 で割り切れないならば、 a * (a ^ (n - 1)) を計算する
  }
}

以上の関数が実装されている状態で、 main() 関数に答えを出力する部分を直接実装しました。

const ll MOD2 = 998244353;

int h, w;
char s[1005][1005];
int group_count = 0;

int main() {
  // 値の入力
  cin >> h >> w;
  for (int i = 1; i <= h; i++) {
    for (int j = 1; j <= w; j++) {
      cin >> s[i][j];
    }
  }

  // すべての未探索の緑色の (i, j) について、dfs を行う
  for (int i = 1; i <= h; i++) {
    for (int j = 1; j <= w; j++) {
      if (group[i][j] == 0 && s[i][j] == '#') {
        group_count++;
        dfs(i, j);
      }
    }
  }

  ll red_count = 0, green_sum = 0; // red_count は赤マスの総数、green_sum は塗り替え後の緑の連結成分数の総数をそれぞれ表す

  // すべての赤マスに対して、緑の連結成分数を計算する
  for (int i = 1; i <= h; i++) {
    for (int j = 1; j <= w; j++) {
      if (s[i][j] == '.') {
        red_count++; // red_count を 1 増やす
        green_sum += search(i, j); // 緑の連結成分数を green_sum に足す
        green_sum %= MOD2;
      }
    }
  }

  cout << green_sum * my_pow(red_count, MOD2 - 2) % MOD2 << endl; // (green_sum / red_count) を mod 998244353 で出力する

  return 0;
}

感想

解いていたときにはシンプルな問題かなと思っていたのですが、実装にて思っていたよりも多くの関数を作ることになっていました。 確率や期待値の問題は、数え上げになにかもう1要素加わっている感があって、個人的には好きです。

*1: 期待値の  \mathrm{mod} については 過去の記事 で紹介しています。