🤖

# キーエンスプログラミングコンテスト2021 C - Robot on Grid

2021/01/26に公開

動的計画法

キーエンスプログラミングコンテスト2021 C - Robot on Grid

問題へのリンク

問題概要

H 行 W 列のマス目がある。それぞれのマスには R, D, X のいずれかの文字を書き込むことができる。
すぬけ君は $K$ 個のマスを選んで文字を書き込んだ。 $i$ 番目に書き込んだマスは $(h_i, w_i)$ で、書き込んだ文字は $c_i$ であった。

残りのマスへの文字の書き込み方は $3^{HW - K}$ あるが、ロボットが (1, 1) から $(H, W)$ へたどり着けるような場合をすべて足し $998244353$ で割ったあまりを求めよ。

$(i, j)$ が R ならばロボットは $(i, j + 1)$ にのみ移動できる
$(i, j)$ が D ならばロボットは $(i + 1, j)$ にのみ移動できる
$(i, j)$ が X ならばロボットは $(i, j + 1)$ , $(i + 1, j)$ のどちらにも移動できる

制約

$2 \leq H, W \leq 5000$
$0 \leq K \leq min(HW, 2 × 10^5)$
$1 \leq h_i \leq H$
$1 \leq w_i \leq W$

ABC中の解答

少し前のコンテストで出てきた縦、横、斜めの累積和を持つDPかなと思って過去問を見たら見つけたのが ABC183 E - Queen on Grid。DPを使うのは間違いなさそうだなと思って配るDPで実装してみたけどsample 2つしか正解できず時間切れ。

H, W, K = map(int, input().split())
MOD = 998244353

cells = {}
for _ in range(K):
    h, w, c = input().split()
    h, w = int(h) - 1, int(w) - 1
    cells[(h, w)] = c
# print(f'{cells=}')

# dp[i][j] : マス (i, j) まで移動する方法の数

dp = [[0] * (H + 1) for _ in range(W + 1)]
dp[0][0] = 1

v1 = pow(3, H * W - K - 1, MOD)
v2 = pow(3, H * W - K, MOD)
v3 = pow(3, (H * W - K) * (H - 1 + W - 1 - 1), MOD)
for i in range(H):
    for j in range(W):
        v = cells.get((i, j), None)
        if v is None:
            # D + X
            dp[i + 1][j] = (dp[i + 1][j] + dp[i][j] * v1 * 2) % MOD
            # R + X
            dp[i][j + 1] = (dp[i][j + 1] + dp[i][j] * v1 * 2) % MOD
        elif v == 'D':
            dp[i + 1][j] = (dp[i + 1][j] + dp[i][j] * v2) % MOD
        elif v == 'R':
            dp[i][j + 1] = (dp[i][j + 1] + dp[i][j] * v2) % MOD
        else:  # X
            dp[i + 1][j] = (dp[i + 1][j] + dp[i][j] * v2) % MOD
            dp[i][j + 1] = (dp[i][j + 1] + dp[i][j] * v2) % MOD
ans = dp[H - 1][W - 1] // v3
print(ans)

公式解法1

各セルに行ける書き込み方の個数を求めて最後 $(i, j)$ に行ける書き込みの個数が答えと等しくなる。また、空白マスは X と書かれたマスと同様に右、下のどちらにも移動可能とみなす。

この時、 $dp(h, w, k)$ := ロボットが $(1, 1)$ から出発して空白マスを $k$ 箇所通って $(h, w)$ に到達するような経路の個数と定義する。この時、通った空白マスは X か　(R, D の移動方向に対応するどちらか)で書き込む必要があり、その他の通ってない空白マスはどの文字を書いても良い。すると、 $\Sigma_{1 \leq k \leq H + W} dp(H, W, k) 2^k 3^{HW - K - k}$ が求める答えになる。

計算量を下げるために $2^k 3^{HW - K -k}$ の処理を最後までしないところもポイントだなと思った。

import numpy as np

H, W, K = map(int, input().split())
MOD = 998244353

cells = {}
for _ in range(K):
    h, w, c = input().split()
    h, w = int(h) - 1, int(w) - 1
    cells[(h, w)] = c

# dp[h][w][k] := ロボットが(0, 0)から出発して空白マスをk箇所通って(h - 1, w - 1)に到達するような経路の個数
dp = np.zeros((H + W + 1, H + 1, W + 1), dtype='i')
dp[0, 0, 0] = 1

# 配るDP
for h in range(H):
    for w in range(W):
        v = cells.get((h, w), None)
        for k in range(H + W):
            if v is None:
                dp[k + 1, h + 1, w] += dp[k, h, w]
                dp[k + 1, h, w + 1] += dp[k, h, w]
            if v in ['D', 'X']:
                dp[k, h + 1, w] += dp[k, h, w]
            if v in ['R', 'X']:
                dp[k, h, w + 1] += dp[k, h, w]

ans = 0
for k in range(H + W + 1):
    ans += dp[k, H - 1, W - 1] * pow(2, k, MOD) * pow(3, H * W - K - k)
    ans %= MOD
print(int(ans))

しかし、これだと計算量が $O(HW(H + W))$ となってしまう。今回の制約では最大 $O(5 \times 10^3 \times 5 \times 10^3 (2 \times 5 \times 10^3))$ となり計算が間に合わない。( $H + W$ は $(1, 1)$ から $(H, W)$ へ行く時に通る空白マスの最大の値であり $H \times W$ より少ないのでこちらを考慮したほうがよい)

そこで、 $k$ を状態にもつ変わりに遷移の際に $\cfrac{2}{3}$ をかけて最後に $3^{HW - K}$ をかければ計算量が $O(HW)$ となり計算が間に合うようになる。(結局上のコードでは $k$ の数で最後に $\cfrac{2}{3}$ を何回かけるかを決めているだけなので空白マスを通るたびに毎回 $\cfrac{2}{3}$ かけるのと同じ作業になる)

と、解説を読んで理解はしたが $\cfrac{2}{3}$ をかけると整数じゃなくなってMODができなくなるのでどうするんだろうと思って maspy さんのコードを見たら MOD (998244353) の世界で $\cfrac{2}{3}$ で割ったことになるように (MOD + MOD + 2) // 3 をかけていた。なるほどなー

というわけでなんとかAC

H, W, K = map(int, input().split())
MOD = 998244353

cells = {}
for _ in range(K):
    h, w, c = input().split()
    h, w = int(h) - 1, int(w) - 1
    cells[(h, w)] = c

# dp[w][k] := ロボットが(0, 0)から出発して(h - 1, w - 1)に到達するような経路の個数, 配る際に2/3をかける
dp = [[0] * (W + 1) for _ in range(H + 1)]
dp[0][0] = 1

c = (MOD + MOD + 2) // 3  # 2/3

# 配るDP
for h in range(H):
    for w in range(W):
        dp[h][w] %= MOD
        v = cells.get((h, w), None)
        if v is None:
            dp[h + 1][w] += dp[h][w] * c
            dp[h][w + 1] += dp[h][w] * c
            continue
        if v in ['D', 'X']:
            dp[h + 1][w] += dp[h][w]
        if v in ['R', 'X']:
            dp[h][w + 1] += dp[h][w]

ans = dp[H - 1][W - 1] * pow(3, H * W - K, MOD)
ans %= MOD
print(ans)

cf. maspyさんのコード