ABC167 E - Colorful Blocks

方針ガチャ、成功！

問題のリンク

E - Colorful Blocks

コンテスト中の思考過程

DP？

・dp[i][j] := 前からi 個目のブロックまで見て、今まで同色で塗られた隣り合うブロックの組がj 個ある場合のブロックの並べ方の数。
・これはO(NK) でTLE。この方針で高速化する方法は思いつかなかったので、やめた。

ちょっと簡単な場合を考えてみる

・K=0、隣り合う同色のブロックの組が一つも無い場合
→先頭の色の選び方はM通り、次はM-1 通り、次もM-1通り....... でこのときの答えは、
・隣り合う同色のブロックがちょうど1組の場合。
→例えば、先頭とその次のブロックを同じ色で塗るとすると、先頭のブロックの色の選び方はM通り。その次のブロックの色は先頭と同じなので1通り。それ以降は全てM-1通り。このときの答えは
→異なる組み合わせのブロックが同色であっても、同様。どのブロックの組を隣り合わせるかは、 通り。 よってこのときの答えは
・隣り合う同色のブロックがちょうどi組の場合。
→同様に考えて、答えは、

・よって、

以下実装

#include <iostream>
#include <algorithm>
#include <vector>
#define rep(i,n) for(int i=0;i<n;++i)
#define rep1(i,n) for(int i=1;i<=n;++i)
using namespace std;
template<class T>bool chmax(T &a, const T &b) { if(a < b){ a = b; return 1; } return 0; }
template<class T>bool chmin(T &a, const T &b) { if(a > b){ a = b; return 1; } return 0; }
typedef long long ll;
const ll MOD=998244353;
struct mint{
  ll x;
  mint(ll x=0):x(x%MOD){}
  mint& operator+=(const mint a){
    if((x+=a.x)>=MOD) x-=MOD;
    return *this;
  }
  mint& operator-=(const mint a){
    if((x += MOD-a.x)>=MOD) x-=MOD;
    return *this;
  }
  mint& operator*=(const mint a){
    (x*=a.x)%=MOD;
    return *this;
  }
  mint operator+(const mint a) const{
    mint res(*this);
    return res+=a;
  }
  mint operator-(const mint a) const{
    mint res(*this);
    return res-=a;
  }
  mint operator*(const mint a) const{
    mint res(*this);
    return res*=a;
  }
  mint pow(ll t) const{
    if(!t) return 1;
    mint a = pow(t>>1);
    a*=a;
    if(t&1) a*=*this;
    return a;
  }
  
  //for prime mod
  mint inv() const{
    return pow(MOD-2);
  }
  mint& operator/=(const mint a){
    return (*this) *= a.inv();
  }
  mint operator/(const mint a) const{
    mint res(*this);
    return res/=a;
  }
};
const int NMAX=200010; // we can calculate nCk until n is equal to NMAX
mint fact[NMAX],infac[NMAX];
void Make_Fact(){
  fact[0]=fact[1]=1;
  for(int i=2;i<=NMAX-1;++i){
    fact[i]=fact[i-1]*(mint)i;
  }
}
void Make_InvFact(){
  infac[0]=infac[1]=1;
  for(int i=2;i<=NMAX-1;++i){
    infac[i]=infac[i-1]/(mint)i;
  }
}
mint Comb(int n,int k){
  if(n<0||k<0||n-k<0) return 0;
  return fact[n]*infac[k]*infac[n-k];
}
mint pow2[200010];

int main()
{
  int n,m;cin >> n >> m;
  int k; cin >> k;

  Make_Fact();
  Make_InvFact();

  pow2[0].x = 1;
  rep(i,n) pow2[i+1] = pow2[i]*(mint)(m-1);

  mint res;

  rep(i,k+1) {
    res += (mint)m * pow2[n-1-i] * Comb(n-1, i);
  }
  
  cout << res.x << "\n";
  return 0;
}