atcoder.jp
問題
- H行W列のグリッドが与えられるよ
- グリッドが赤、緑の2色で塗り分けられてるよ
- 赤色のマスをランダムで1つ選んで、緑色に塗り替えるよ
- 塗替えた後の、緑色の連結成分数の期待値をmod998244353でお願い!
成約
- $ 1 \leq H,W \leq 1000$
- 赤色のマスは絶対に存在する
思考
- R=赤色のマスの数の時、1/Rの確率で、赤色のマスが選ばれるね
- 期待値 = 「結果×その結果が起きる確率、の総和」
- その結果が起きる確率は常に1/Rなので、
- 期待値 = 「結果の総和」×1/Rで良さそう
- じゃあ結果の話なんですけど、まずは連結成分をUnionFindでそれぞれ特定してあげて…
- 赤色マスの、上下左右にいくつグループが存在するか?で、それぞれグループ数の変動を数えて上げれば良さそう
コード
from atcoder.dsu import DSU
H,W = map(int,input().split())
DSU = DSU(H*W)
grid_L = []
for i in range(H):
grid_L.append(input())
dir = [(1, 0),(-1, 0),(0, 1),(0, -1)]
green_cnt = 0
red_cnt = 0
for r in range(H):
for c in range(W):
if grid_L[r][c] == ".":
red_cnt += 1
continue
green_cnt += 1
for dr,dc in dir:
nr,nc = r+dr,c+dc
if 0 <= nr < H and 0 <= nc < W and grid_L[nr][nc] == "#":
DSU.merge(r*W+c,nr*W+nc)
ans = 0
l = [1,0,-1,-2,-3]
group_cnt = 0
for r in range(H):
for c in range(W):
if grid_L[r][c] == "#":
if DSU.leader(r*W+c) == r*W+c:
group_cnt += 1
else:
tmp_s = set()
for dr,dc in dir:
new_r = r + dr
new_c = c+dc
if 0 <= new_r < H and 0 <= new_c < W and grid_L[new_r][new_c] == "#":
tmp_s.add(DSU.leader(new_r*W+new_c))
ans += l[len(tmp_s)]
ans += red_cnt * group_cnt
mod = 998244353
inv = pow(red_cnt,-1,mod)
print(ans*inv%mod)