yukicoder : No.1897 Sum of 2nd Max

ちょっと手間取った。
https://yukicoder.me/problems/no/1897

問題

1～Kの整数値からなるN要素の整数列はK^N通りある。
各数列において、2番目に大きな値の総和を答えよ。

解法

f(x) := K^N通りのうち、2番目に大きな値がx以上である組み合わせ数
とすると、
g(x) := K^N通りのうち、2番目に大きな値がちょうどxである組み合わせ数
はg(x)=f(x)-f(x-1)であり、また解はsum(x*g(x))である。

あとはf(x)を求めよう。
f(x)は「N要素中x以上の値が2個以上ある組み合わせ」なので、言い換えると「K^N通りから、x以上が0個のケースと1個のケースを除いたもの」で計算できる。

int N,K;
const ll mo=998244353;
ll dp[202020];

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>>K;
	ll ret=0;
	for(i=K;i>=1;i--) {
		dp[i]=modpow(K,N);
		// i以上が1個だけ
		dp[i]-=1LL*(K-i+1)*N%mo*modpow(i-1,N-1);
		// i以上が0個だけ
		dp[i]-=modpow(i-1,N);
		dp[i]=(dp[i]%mo+mo)-mo;
		
		ret+=i*(mo+dp[i]-dp[i+1])%mo;
	}
	cout<<ret%mo<<endl;
}

まとめ

結構悩んだけど、解いてみればさほど難しくない。