くふむ

けふぇめ

マスターオブ場合の数を競プロで殴る その1「サイコロの積」

私の大きな弱点である数え上げを克服するために、『大学への数学』で有名な東京出版の『マスター・オブ・場合の数』第1章を解いていくことにしました。

もちろん紙とペンを使って解くこともしますが、競プロに最大限生かすため、手で解くだけに留めずプログラムでもぶん殴ってみよう！ という遊びを兼ね備えた精進、精進を兼ね備えた遊びです。

ただし、『マスター・オブ・場合の数』の問題は手で解くことを想定しているため、そのままプログラミングで殴ると多くの問題が全探索で通ってしまいます。そこで、問題を一般化して競プロっぽく改題し、それに対してできる限り大きい制約で解けるようチャレンジしていこうと思います。

「ここがおかしい！」や「もっと計算量落とせるよ！」などのご指摘は歓迎です！

問題

$1$ から $6$ の目が出るサイコロを $N$ 回振るとき、目の出方は $6^{N}$ 通りあります。このうち、出る目の積が $2^{M}$ の倍数となるような目の出方は何通りありますか？

答えを $998244353$ で割ったあまりを求めてください。

制約

  • $1 \leq N \leq 3000$
  • $1 \leq M \leq 3000$

『マスター・オブ・場合の数』P.12 §1-1. §1-2. をもとに改題

おそらく水diffぐらいの難易度だと思います。

クリックで解答を表示

解答

「出た目の積が $2^{M}$ の倍数」という条件は「 $2$ の素因数の個数が ${M}$ 個以上」と言い換えられます。

$1 \sim 6$ を素因数分解して、$2$ の素因数の個数を調べてみましょう。

  • $1 = 2^{\mathbf{0}}$
  • $2 = 2^{\mathbf{1}}$
  • $3 = 2^{\mathbf{0}} \cdot 3^{1}$
  • $4 = 2^{\mathbf{2}}$
  • $5 = 2^{\mathbf{0}} \cdot 5^{1}$
  • $6 = 2^{\mathbf{1}} \cdot 3^{1}$

$2$ の素因数の個数は $1, 3, 5$ が出たら $0$ 個、$2, 6$ が出たら $1$ 個、$4$ が出たら $2$ 個増えるということが分かりました。これを遷移とする以下のようなDPを行うことで、この問題を解くことができます。

DP配列の定義

$dp[i][j] =$ サイコロを $i$ 回ふって $2$ の素因数の個数がちょうど $j$ 個であるような場合の数 $\pmod{998244353}$

ただし $j \gt {M}$ の場合は $j={M}$ にまとめて扱うことにします。

初期値

  • $dp[0][0] = 1$
  • $dp[i][j] = 0\hskip2em (otherwise)$

遷移

$j \geq {M}$ の場合をまとめて扱う都合上、配るDPのほうが直感的に書けます。

  • $dp[i + 1][j] \mathrel{{+}{=}} dp[i][j] \times 3$
  • $dp[i + 1][\min(j+1, M)] \mathrel{{+}{=}} dp[i][j] \times 2$
  • $dp[i + 1][\min(j+2, M)] \mathrel{{+}{=}} dp[i][j] \times 1$

貰うDPで書く場合は $j={M}$ のときの場合分けに注意が必要です。

  • $dp[i + 1][j] = dp[i][j] \times 3 + dp[i][j - 1] \times 2 + dp[i][j - 2] \times 1 \hskip1em (j \lt M)$
  • $dp[i + 1][j] = dp[i][j] \times 6 + dp[i][j - 1] \times 3 + dp[i][j - 2] \times 1 \hskip1em (j = M)$

更新順

小さい方からでOK。

このDP配列を埋めたとき、$dp[N][M]$ が求める答えです。計算量は $O(NM)$ です。

実装例

配るDPで実装した例です。

#include <iostream>
using namespace std;
constexpr int MOD = 998244353;

int dp[3001][3001];
int main() {
    // input
    int n, m;
    cin >> n >> m;

    // dp 配列の初期化
    for (int i = 0; i <= n; ++i)
        for (int j = 0; j <= m; ++j)
            dp[i][j] = 0;
    dp[0][0] = 1;
        
    // dp 遷移
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j <= m; ++j) {
            dp[i + 1][j] += dp[i][j] * 3;
            dp[i + 1][min(j + 1, m)] += dp[i][j] * 2;
            dp[i + 1][min(j + 2, m)] += dp[i][j];

            dp[i + 1][j] %= MOD;
        }
    }

    if (dp[n][m] < 0) dp[n][m] += MOD;

    // output
    cout << dp[n][m] << endl;
}

貰うDPで書く場合は、DP遷移の部分が次のように変わります（他は同じです）。

    // dp 遷移
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j <= m; ++j) {
            if (j < m)
                dp[i + 1][j] += dp[i][j] * 3 + dp[i][j - 1] * 2 + dp[i][j - 2];
            else
                dp[i + 1][j] += dp[i][j] * 6 + dp[i][j - 1] * 3 + dp[i][j - 2];

            dp[i + 1][j] %= MOD;
        }
    }

余談

『マスター・オブ・場合の数』には、多項式の次数に置き換えて解く方法が紹介されていました。なぜこれが成り立つのかについては本を読んでいただくとして、この問題は次の問題に帰着できます。

$\sum^{\infty}_{k=M}[x^{k}](x^{2}+2x+3)^{N}$ を $998244353$ で割ったあまりを求めよ。 ただし $[x^{k}]f(x)$ は $f(x)$ の $k$ 次の係数を表す。

これを見ると繰り返し2乗法で効率的に解けそうな気がしてくるのですが、シグマを何とかする方法が思いつきませんでした。何かいい方法はあるのでしょうか。分かる人がいたら是非教えて下さい。

類題： TDPC D - サイコロ

サイコロを $N$ 回振ったとき、出た目の積が $D$ の倍数となる確率を求めよ。

制約

  • $1 \leq N \leq 100$
  • $1 \leq D \leq 10^{18}$

https://atcoder.jp/contests/tdpc/tasks/tdpc_dice

TDPCに似てる問題があったのでこちらも解いてみます。

クリックで解答を表示

解答

先程とほとんど同じ問題設定ですが、$2^{M}$ の倍数から $D$ の倍数へと条件が広がっています。$D$ を素因数分解できればさっきと同じ手順で解くことができるはずですが、 $1 \leq D \leq 10^{18}$ という制約では素因数分解ができません（素因数分解には通常 $O(\sqrt{D})$ かかるため）。

仕方ないので小さいケースで色々試していると、$D = 7$ のとき $N$ によらず答えが $0$ になることに気付きました。$D=11$ などでも同様です。

これを踏まえてよく考えてみると、$1 \sim 6$ に含まれる素因数は $2, 3, 5$ しかありえません。そのため、$\boldsymbol{D}$ に $\mathbf{2, 3, 5}$ 以外の素因数が含まれていたらその時点で確率は $\mathbf{0}$ になるのです。素因数が $2, 3, 5$ しかないと分かれば、ひたすらこれらの数で割っていくことで $O(\log{D})$ で素因数分解ができます。

$D=2^{A} \cdot 3^{B} \cdot 5^{C}$ の形に素因数分解できたら、あとは先程と同様です。

  • $dp[i][a][b][c] =$ サイコロを $i$ 回ふって $2, 3, 5$ の素因数の個数がそれぞれ $a, b, c$ 個であるような確率 $\pmod{998244353}$

  • $dp[0][0][0][0] = 1$

  • $dp[i][a][b][c] = 0\hskip2em (otherwise)$

と定義し、$1 \sim 6$ が

  • $1 = 2^{\mathbf{0}} \cdot 3^{\mathbf{0}} \cdot 5^{\mathbf{0}}$
  • $2 = 2^{\mathbf{1}} \cdot 3^{\mathbf{0}} \cdot 5^{\mathbf{0}}$
  • $3 = 2^{\mathbf{0}} \cdot 3^{\mathbf{1}} \cdot 5^{\mathbf{0}}$
  • $4 = 2^{\mathbf{2}} \cdot 3^{\mathbf{0}} \cdot 5^{\mathbf{0}}$
  • $5 = 2^{\mathbf{0}} \cdot 3^{\mathbf{0}} \cdot 5^{\mathbf{1}}$
  • $6 = 2^{\mathbf{1}} \cdot 3^{\mathbf{1}} \cdot 5^{\mathbf{0}}$

と素因数分解されることから、DPの遷移は

  • $dp[i+1][\min(a+0 ,A)][\min(b+0 ,B)][\min(c+0 ,C)] \mathrel{{+}{=}} dp[i][a][b][c] / 6$
  • $dp[i+1][\min(a+1 ,A)][\min(b+0 ,B)][\min(c+0 ,C)] \mathrel{{+}{=}} dp[i][a][b][c] / 6$
  • $dp[i+1][\min(a+0 ,A)][\min(b+1 ,B)][\min(c+0 ,C)] \mathrel{{+}{=}} dp[i][a][b][c] / 6$
  • $dp[i+1][\min(a+2 ,A)][\min(b+0 ,B)][\min(c+0 ,C)] \mathrel{{+}{=}} dp[i][a][b][c] / 6$
  • $dp[i+1][\min(a+0 ,A)][\min(b+0 ,B)][\min(c+1 ,C)] \mathrel{{+}{=}} dp[i][a][b][c] / 6$
  • $dp[i+1][\min(a+1 ,A)][\min(b+1 ,B)][\min(c+0 ,C)] \mathrel{{+}{=}} dp[i][a][b][c] / 6$

と書けます。添字が小さい順に埋めていって、$dp[N][A][B][C]$ が求める答えとなります。計算量は $O(NABC)$ つまり $O(N\log^{3}{D})$ となります。

ACコード(Python, 141ms)