黒ココアさんのメモ置き場

メモを置いたり置かなかったり

ABC334 E - Christmas Color Grid 1

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)