AtCoder ARC #147 : D - Sets Scores

コードが短い。
https://atcoder.jp/contests/arc147/tasks/arc147_d

問題

整数N,Mが与えられる。
N個の整数集合S[i]からなる列Sを考える。

• S[i]は1～Mの範囲の整数の集合である
• S[i]とS[i+1]のうち、片方にのみ含まれる要素は1個である。

f(k) := S[0]～S[N-1]のうち、kを含む個数
としたとき、すべてのSのパターンにおけるf(1)*f(2)*....*f(M)の総和を求めよ。

解法

S[i]とS[i+1]で有無が変わる要素は1個であり、各iにおける全パターンはM^(N-1)通りである。
変化のパターンを1個決めると、S[0]が決まれば各Sが決まる。対称性から、この時のf(1)*f(2)...はN^Mとなる。
よって解はM^(N-1)*N^Mとなる。

int N,M;
const ll mo=998244353;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	cin>>N>>M;
	cout<<modpow(M,N-1)*modpow(N,M)%mo<<endl;
}

まとめ

考察・実験は手間取ったけどコードは短いパターン。