Keyence2021 C - Robot on Grid解説[python]

公開:2021/01/17
更新:2021/01/17
3 min読了の目安(約2800字TECH技術記事

URL

https://atcoder.jp/contests/keyence2021/tasks/keyence2021_c

制約

2 <= H,W<=5000
0<= K <= min(HW,2*10^5)
1 <= hi<= H
1 <= wi<= W

問題概要

  • H*Wのgrid がある。それぞれのマスにはR,D,Xのいずれかが書き込める
  • Kこのマス(hi,wi)を選んで文字Ciを書き込んだ
  • 残りのHW-K のそれぞれの書き込み方(3^(HW-K)通り)での移動方法の総和を998244353で割ったあまりを答えよ
  • (1,1)から(H,W)まで移動できるようなルート
  • R ならj+1 , Dならi+1に移動。X なら両方いける

提出コード

def xgcd(a, b):
    x0, y0, x1, y1 = 1, 0, 0, 1
    while b != 0:
        q, a, b = a // b, b, a % b
        x0, x1 = x1, x0 - q * x1
        y0, y1 = y1, y0 - q * y1
    return a, x0, y0

def modinv(a, m):
    g, x, y = xgcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m

h, w, k = map(int, input().split())
dp = [[0] * (w + 2) for _ in range(h + 2)]
grid = [[0] * (w + 2) for _ in range(h + 2)]
for i in range(k):
    p, q, r = map(str, input().split())
    p = int(p)
    q = int(q)
    if r == 'X': grid[p][q] = 1
    if r == 'R': grid[p][q] = 2
    if r == 'D': grid[p][q] = 3
mod = 998244353
inv3 = modinv(3,mod)

dp[1][1] = 1
for i in range(1, h + 1):
    for j in range(1, w + 1):
        if grid[i][j] == 1:
            dp[i + 1][j] += dp[i][j]
            dp[i][j + 1] += dp[i][j]
        elif grid[i][j] == 2:
            dp[i][j + 1] += dp[i][j]
        elif grid[i][j] == 3:
            dp[i + 1][j] += dp[i][j]
        else:  # 空白文字
            dp[i + 1][j] += 2 * inv3 * dp[i][j]
            dp[i][j + 1] += 2 * inv3 * dp[i][j]
        dp[i + 1][j] %= mod
        dp[i][j + 1] %= mod
# print(dp)
print(int(dp[h][w] * pow(3, h * w - k, mod)) % mod)

考察

  • 問題から全探索はできないので、同じような状態をまとめる必要がある。
  • それぞれのパスに対して、それを使ってゴールする書き込み方が何通りあるかを数えたい
    • 書き込み方はそのパス上にある書き込まれていないマスをi個とした時、
      2^{i}*3^{HW-i-K}
      となる
    • 証明:パス上のマスはXかRorDのうちパスに従うものの2通りあり、パス上でないマスはどの文字でも良いため
  • 陽にパスに対して全探索すると1000C500通りでTLEするのでDPでまとめる必要がある
    • dp[i][j] := (i,j)に到達するような書き込み方でのパスの総和 としたい
    • 一方で、それぞれのpathに対して何個の書き込まれていないマスを通ったかの情報を持ちたい
    • しかし、それを状態としてもつとdp[i][j][k] となって、オーダーが
      5000*5000*5000
      でTLEする。
      • 書き込まれていないマスを通ったらパスの総和を3で割ることで帳尻を合わせる
    • 答えはdp[h][w]* 3^(HW-K) となる。
  • 漸化式を考えると、配るDPの方が元のマスの値を計算しやすいので、それで考えて
    • if grid[i][j] == "X" : dp[i+1][j] += dp[i][j],dp[i][j+1] += dp[i][j]
    • elif grid[i][j] == "R": dp[i][j+1] += dp[i][j]
    • elif grid[i][j] == "D" : dp[i+1][j] += dp[i][j]
    • else: p[i+1][j] += 2/3dp[i][j],dp[i][j+1] += 2/3dp[i][j]
      という式になる

実装方針

  • 3で割り続けるので、初期値をどうするかがきも
    • modinv 逆元を使う。

次回への反省

参考

https://tex2e.github.io/blog/crypto/modular-mul-inverse 
python3でのモジュラ逆数の求め方