aruma256

arumaです

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 のときの初期状態は下図のように表現できます。

S
S
G
G
1
1
0
0
0
0
0
0
0
0
step 0
step 0Text is not SVG - cannot display

step 0 → step 1

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

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

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

1
1
0
0
0
0
0
0
0
0
0
0
1/3
1/3
1/3
1/3
1/3
1/3
0
0
出目 : ①
1×1/3
出目 : ①...
③ 1×1/3
③ 1×1/3
②
1×1/3
② 1×1/3
step 0
step 0
step 1
step 1Text is not SVG - cannot display

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

step 1 → step 2

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

0
0
1/3
1/3
1/3
1/3
1/3
1/3
0
0
0
0
0
0
1/9
1/9
5/9
5/9
3/9
3/9
1/3×1/3
1/3×1/3
1/3×1/3 + (1/3×1/3) ×2 + (1/3×1/3) ×2
1/3×1/3 + (1/3×1/3) ×2 + (1/3×1/3)...
step 1
step 1
step 2
step 2Text is not SVG - cannot display

step 2 → step 3

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

0
0
0
0
1/9
1/9
5/9
5/9
3/9
3/9
0
0
0
0
5/27
5/27
7/27
7/27
15/27
15/27
1/9×1/3 + 5/9×1/3 + 3/9
1/9×1/3 + 5/9×1/3 + 3/9
step 2
step 2
step 3
step 3Text is not SVG - cannot display

ここまでをまとめる

この操作を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, … ] は前から順に高速に求められるので、これを事前に計算しておき、それらを組み合わせて目的の値を得るということが可能です。

解法

以上より、

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

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

Python標準のpow関数

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

回答例

You May Also Enjoy