ABC300-E Dice Product 3
/icons/hr.icon
出典
/icons/hr.icon
問題
あなたは$ 1以上 $ 6以下の整数が等確率で出るサイコロと整数$ 1を持っています。
あなたは持っている整数が$ N未満である間、次の操作を繰り返します。
サイコロを振り、出た目を$ xとする。持っている整数に$ xを掛ける。
全ての操作を終了した時に、持っている整数が$ Nに一致する確率を$ \mod 998244353で求めてください。
/icons/hr.icon
制約
$ 2 \leq N \leq 2 \times 10^{18}
/icons/hr.icon
解法
メモ化再帰を用いて期待値 dp で解くことを考えます。
$ {\mathrm dp} [i] := 目が i になる確率 とします。
$ {\mathrm dp}[1] = 1 です。
$ {\mathrm dp}[i] = \frac{1}{6} \left( {\mathrm dp}[i] + {\mathrm dp}[2i] + {\mathrm dp}[3i] + {\mathrm dp}[4i] + {\mathrm dp}[5i] + {\mathrm dp}[6i] \right)
$ \frac{5}{6} {\mathrm dp}[i] = \frac{1}{6} \left( {\mathrm dp}[2i] + {\mathrm dp}[3i] + {\mathrm dp}[4i] + {\mathrm dp}[5i] + {\mathrm dp}[6i] \right)
$ {\mathrm dp}[i] = \frac{1}{5} \left( {\mathrm dp}[2i] + {\mathrm dp}[3i] + {\mathrm dp}[4i] + {\mathrm dp}[5i] + {\mathrm dp}[6i] \right)
このように式変形を行うことで解くことができました。
/icons/hr.icon
コード
code:python
from collections import defaultdict
from heapq import heappush, heappop, heapify
def main() -> None:
N = int(input())
dp = defaultdict(lambda: 0)
mod = 998244353
heapify(q)
s = set()
s.add(1)
while q:
v = heappop(q)
for i in range(2, 7):
if v * i > N:
break
dpv * i += dpv * pow(5, -1, mod) if v * i not in s:
s.add(v * i)
heappush(q, v * i)
if __name__ == "__main__":
main()