Chương trình C++ đếm số cách tô màu khối với 2 điều kiện

Cho N, M, K. Có N khối xếp hàng ngang. Tô mỗi khối một trong M màu (không bắt buộc dùng hết). Tối đa K cặp khối liền kề cùng màu. Đếm số cách tô, modulo 998244353.

Ví dụ

Với N = 3; M = 2; K = 1, đầu ra 6. Các cách: 112, 121, 122, 211, 212, 221.

Giải thuật

Tiền xử lý giai thừa và nghịch đảo modulo. Duyệt i từ 0 đến K (số cặp liền kề cùng màu). Với i cặp cùng màu, có C(N-1, i) cách chọn vị trí các cặp. Màu đầu tiên có M lựa chọn, các đoạn còn lại mỗi đoạn M-1 lựa chọn. Tổng: sum_{i=0}^K C(N-1,i) * M * (M-1)^{N-i-1}.

Mã 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, M = 2, K = 1;
    cout << solve(N, M, K) << endl;
}

Đầu vào / Đầu ra

Input: 3, 2, 1
Output: 6