ABC275Eを解いた【動的計画法,mod998244353】

筆者はレート800前後の茶～緑コーダ

ABC275のE問題を解いていく

実装コード

この問題のポイントはmod998244353におけるMの逆数?(以下iM)を求めることらしい。
Python3.8以降のpow関数は剰余の計算に対応しているので、-1乗すればOK

Mの逆元を求める例

mod = 998244353
inv = lambda x: pow(x, -1, mod)
iM = inv(M)

逆元を求めたあとは素直に動的計画法をすればいいようだ。

• dp[i][j]: i回コマを動かしたあとjマスにいる確率
• dp[0][0] = 1 : 初期値
• dp[i+1][j+m] = dp[i][j]*iM: jマスにいて、i回目のルーレットでmが出たときの遷移

main.py

import sys
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
def err(*args): print(*args, file=sys.stderr)

class ModInt:
    def __init__(self, x, mod = 998244353):
        self.mod = mod
        self.x = x.x if isinstance(x, ModInt) else x % self.mod
    __str__ = lambda self:str(self.x)
    __repr__ = __str__
    __int__ = lambda self: self.x
    __index__ = __int__
    __add__ = lambda self, other: ModInt(self.x + ModInt(other).x)
    __sub__ = lambda self, other: ModInt(self.x - ModInt(other).x)
    __mul__ = lambda self, other: ModInt(self.x * ModInt(other).x)
    __pow__ = lambda self, other: ModInt(pow(self.x, ModInt(other).x, self.mod))
    __truediv__ = lambda self, other: ModInt(self.x * pow(ModInt(other).x, self.mod - 2, self.mod))
    __floordiv__ = lambda self, other: ModInt(self.x // ModInt(other).x)
    __lt__ = lambda self, other: self.x < ModInt(other).x
    __gt__ = lambda self, other: self.x > ModInt(other).x
    __le__ = lambda self, other: self.x <= ModInt(other).x
    __ge__ = lambda self, other: self.x >= ModInt(other).x
    __eq__ = lambda self, other: self.x == ModInt(other).x
    __ne__ = lambda self, other: self.x!= ModInt(other).x

class mint(ModInt):
    pass

mod = 998244353
inv = lambda x: pow(x, -1, mod)

def main():
    N, M, K = rLI()
    
    iM = inv(M)
    
    dp = [[mint(0)]*-~N for _ in range(K+1)]
    
    dp[0][0] = 1
    
    for i in range(K):
        for j0 in range(N):
            for m in range(M):
                j=j0+m+1
                if j > N:
                    j = 2*N-j
                dp[i+1][j]+=dp[i][j0]*iM
        dp[i+1][N]+=dp[i][N]
    ans = dp[K][N]
    print(ans)
    
if __name__ == '__main__':
    main()

まとめ

ABC275Eを解くにあたり、動的計画法を使用した。

期待値をmod998244353で出力するという方式は最初よくわからなかったが、一旦イメージを掴めた気がするので今後似たような問題が出ても萎縮せずに解きたい。