ABC275 E問題の考え方

2022-12-08

ABC275 E - Sugoroku 4 の考え方を書きました。

はじめに

こんにちは。@aruma256 です。
GMOペパボエンジニア Advent Calendar 2022 (第2会場) の8日目の記事です。

私は社内の競技プログラミング好きが集まる #kyopro チャンネルに生息しているので、最近解いた競プロ問題の解説記事を書くことにしました。

競技プログラミングは「提示された要求を満たすコードをいかに早く提出できるか」を競うものです。文字列操作やテキストの入出力など基本的な処理実装はもちろん、難しい問題では高度なアルゴリズムの知識や思考、実装力が要求されます。

今回は AtCoder Beginner Contest 275 E問題の解説です。

より易しい問題の解説記事も投稿しています。「競プロを初めて知った」という方は↓を先にご覧頂くと良いかもしれません。

問題の概要

ABC275 E - Sugoroku 4

双六において、指定回数以内にゴールできる確率を求める
ゴールを超過する出目の場合、折り返す
答えは mod 998244353 で整数として出力する
- 998244353 は素数です（伏線）

有理数表現の確率で考えてみる

mod 998244353 については後で考えることにして、まずは普通の確率計算問題として考えてみます。

step 0

初期状態（ルーレットを “0回” まわしたとき）では、

はじめのマスにいる確率は 1
その他のマスにいる確率は 0

です。

N=5 のときの初期状態は下図のように表現できます。

step 0 → step 1

上の状態からルーレットを1回まわして、出目の数だけコマを進めます。
M = 3 とすると、

1マス進む確率 : 1/3
2マス進む確率 : 1/3
3マス進む確率 : 1/3

なので、各マスにいる確率は下のように変化します。

1ステップごとに、他のマスへ確率が分配されていくイメージです。

step 1 → step 2

ゴールを超過する場合に折り返すというルールのため、異なる出目でも実質的に同じ移動となる場合があります。

step 2 → step 3

ゴールマスの確率は分配や 1/M などせず、単に次stepにて加算されるのみであることに気をつけます。

ここまでをまとめる

この操作をK回繰り返したときのゴールマスに存在している確率が答えです。

このシミュレーションの計算量を見積もると

K回、ルーレットを回す
- N個のマス、それぞれについて
  - M通りの行き先の確率を加算する

より、 O(KNM) です。

確率を mod 998244353 で扱う

さて、先送りにしていた mod 998244353 で確率を表すという部分について考えていきます。

簡単に説明すると、ある素数 p を法とするとき「aで割る」という操作は「モジュラ逆数 a^(p-2) をかける」として扱うことができるので、これを利用します。

参考：フェルマーの小定理 a^(p-1) ≡ 1 (mod p)

モジュラ逆数を高速に求める

a^p-2 を求める計算アルゴリズムにも工夫が必要です。

以下のような単純な実装では、inv_a が非常に大きい数になってしまいます。

# ruby

inv_a = a

(p-3).times {
  inv_a *= a
}

inv_a = inv_a % 998244353

逐次、剰余をとりましょう。

# ruby

inv_a = a

(p-3).times {
  inv_a = (inv_a * a) % 998244353
}

pが小さい数であればこれで十分ですが、pが大きい（目安としては、10^7 より大きい)場合にはさらに高速化が必要です。

繰り返し二乗法

例えば n^104 を求めるとき、 n * n * n * … と103回の掛け算をするよりも効率が良いような計算方法を考えます。

目的の数を2の累乗の和として表現してみると、

104
= 2^3 + 2^5 + 2^6
= 8 + 32 + 64

つまり、

n^104
= n^(8 + 32 + 64)
= n^8 * n^32 * n^64

です。

[n^1, n^2, n^4, n^8, n^16, n^32, n^64, … ] は前から順に高速に求められるので、これを事前に計算しておき、それらを組み合わせて目的の値を得るということが可能です。

解法

以上より、

入力を受け取る。
Mのモジュラー逆数を求める。
以下のシミュレーションを行う
- K回、ルーレットを回す
  - N個のマス、それぞれについて
    - M通りの行き先の確率を加算する
最後のマスに存在している確率が答え

という手順で解くことができます。

Python標準のpow関数

Python3.8以降では、組み込み関数のpowで簡単にモジュラー逆数を求められます。

回答例

Python (3.8.2)
- https://atcoder.jp/contests/abc275/submissions/36069352

Twitter Facebook LinkedIn

aruma256