2つの条件を満たすブロックの塗り方を数えるC++プログラム

3つの整数 N、M、K が与えられます。ここに、一列に並んだ N 個のブロックがあります。次の2つのルールに従ってブロックを塗る方法を考えます。2つの塗り方が「異なる」とみなされるのは、対応するブロックの色が1つでも異なる場合です。

• 各ブロックには、M 色の中から1色を選んで塗ります(すべての色を使い切る必要はありません)。

• 同じ色で塗られた隣接ブロックの組は、最大で K 組まで許容されます。

答えが非常に大きくなる可能性があるため、結果は 998244353 で割った余りを返します。

たとえば、入力が N = 3、M = 2、K = 1 の場合、出力は 6 になります。「112」「121」「122」「211」「212」「221」の6通りの塗り方が条件を満たすためです。

考え方

この問題は、「隣接するブロックが同じ色になる境界」に着目すると整理できます。隣接する境界は全部で N−1 箇所あり、そのうちちょうど i 箇所が同じ色になると仮定すると、塗り方は次のように数えられます。

• 同じ色になる境界を N−1 箇所から i 箇所選ぶ方法は C(N−1, i) 通り。

• ブロック列はこの選び方によって N−i 個のグループに分割されます。各グループ内のブロックはすべて同色で、隣り合うグループは必ず異なる色になります。左端のグループの色は M 通り、以降の各グループは直前と同じ色にできないため M−1 通りずつあるので、全体で M × (M−1)^N−i−1 通りになります。

よって、i を 0 から K まで動かしながら C(N−1, i) × M × (M−1)^N−i−1 を足し合わせれば答えが得られます。二項係数と累乗を毎回計算すると非効率なため、階乗のテーブル fac とそのモジュラ逆元 inv を前計算しておき、O(1) で参照できるようにします。

手順

以上を踏まえると、この問題は次の手順で解けます。

maxm := 2×10^6 + 5
p := 998244353
サイズ maxm の配列 fac と inv を用意する

関数 ppow(a, b, p) を定義する(繰り返し二乗法による累乗計算):
    ans := 1 mod p
    a := a mod p
    b が 0 でない間、次を繰り返す:
        b が奇数ならば:
            ans := ans * a mod p
        a := a * a mod p
        b := b / 2
    ans を返す

関数 C(n, m) を定義する(二項係数の計算):
    m < 0 または m > n の場合:
        0 を返す
    fac[n] * inv[m] mod p * inv[n - m] mod p を返す

main メソッドでは次を実行する:
fac[0] := 1
i := 1 から開始し、i < maxm の間 i を1ずつ増やす:
    fac[i] := fac[i - 1] * i mod p
inv[maxm - 1] := ppow(fac[maxm - 1], p - 2, p)
i := maxm - 2 から開始し、i >= 0 の間 i を1ずつ減らす:
    inv[i] := (i + 1) * inv[i + 1] mod p
ans := 0
i := 0 から開始し、i <= k の間 i を1ずつ増やす:
    t := C(n - 1, i)
    tt := m * ppow(m - 1, n - i - 1, p)
    ans := (ans + t * tt mod p) mod p
ans を返す

実装例

理解を深めるために、以下の C++ 実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;

const long maxm = 2e6 + 5;
const long p = 998244353;
long fac[maxm], inv[maxm];

long ppow(long a, long b, long p){
    long ans = 1 % p;
    a %= p;
    while (b){
        if (b & 1)
            ans = ans * a % p;
        a = a * a % p;
        b >>= 1;
    }
    return ans;
}
long C(long n, long m){
    if (m < 0 || m > n)
        return 0;
    return fac[n] * inv[m] % p * inv[n - m] % p;
}
long solve(long n, long m, long k){
    fac[0] = 1;
    for (long i = 1; i < maxm; i++)
        fac[i] = fac[i - 1] * i % p;
    inv[maxm - 1] = ppow(fac[maxm - 1], p - 2, p);
    for (long i = maxm - 2; i >= 0; i--)
        inv[i] = (i + 1) * inv[i + 1] % p;
    long ans = 0;
    for (long i = 0; i <= k; i++){
        long t = C(n - 1, i);
        long tt = m * ppow(m - 1, n - i - 1, p) % p;
        ans = (ans + t * tt % p) % p;
    }
    return ans;
}
int main(){
    int N = 3;
    int M = 2;
    int K = 1;
    cout << solve(N, M, K) << endl;
}

入力

3, 2, 1

出力

6