E - Colorful Blocks

2023年は数え上げとか期待値とか苦手分野をやっていきます。

問題

以下の条件を満たしながら個のブロックを個の色での塗り方が何通り有るか求めよ。

• 隣り合うブロックの組であって同じ色で塗られている組は、 組以下である。

atcoder.jp

解法

条件の部分を言い換えると、隣接している」かつ「同色」を満たすブロックの組は K 組以下であると言い換えると良いらしい。
日本語が読めないので競プロサークルに助けを求めたら言い換えてくれました。Gくんありがとう。

さて、条件の部分から考察を考える。一度に組以下という条件を扱うのは厄介そうなので、何組使うかを決め打ちして考えることにしよう。(制約的にも間に合う, )
重要な考察として、個のブロックの中で、「隣接するブロックの組」は 個存在する。したがって、同じ色で塗られている組組のとき、個の「隣接するブロックの組」から個を選べば良い。

色の塗り方は、隣接する色と違う色で塗らなければならない。これは、最初だけ色選べて、それ以降は色から選べるというふうに考えればよい。

提出コード

#include <bits/stdc++.h>

using namespace std;

using Int  = long long;
constexpr Int mod = 998244353;

template <class T> T modpow(T a, T b) {
    if (b == 0) return T(1);
    if (b % 2 == 0) {
        long long d = modpow(a, b / T(2));
        return (d * d) % mod;
    } else {
        return (a * modpow(a, b - 1) % mod );
    }
}

template <class T>
class ModComb {
public:
    T maxn,mod;
    vector<T>fac,finv,inv;
    ModComb(T maxn_,T mod_) : maxn(maxn_),mod(mod_),fac(maxn_),finv(maxn_),inv(maxn_){
        fac[0] = fac[1] = 1;
        finv[0] = finv[1] = 1;
        inv[1] = 1;
        for (T i = T(2); i < maxn; i++){
            fac[i] = fac[i - 1] * i % mod;
            inv[i] = mod - inv[mod % i] * (mod / i) % mod;
            finv[i] = finv[i - 1] * inv[i] % mod;
        }
    };
    
    T computeComb(T n,T k){
        if (n < k) return T(0);
        if (n < 0 || k < 0) return T(0);
        return fac[n] * (finv[k] * finv[n - k] % mod) % mod;
    }
    
};

int main()
{
    Int N, M, K;
    cin >> N >> M >> K;
    
    Int ans = modpow(M - 1, N - 1);
    ans *= M; ans %= mod;
    
    auto Comb = ModComb<Int>(1000000LL,mod);
    
    for(Int k = 1; k <= K; ++k)
    {
        Int res = Comb.computeComb(N - 1,k);
        Int a = modpow(M - 1, N - 1 - k);
        a *= M;
        a %= mod;
        a *= res;
        a %= mod;
        ans += a; ans %= mod;
    }
    cout << ans << endl;
}