AtCoder ARC #139 : D - Priority Queue 2

何とか本番中に解けた。
https://atcoder.jp/contests/arc139/tasks/arc139_d

問題

1～Mの値を取るN要素の整数列Aが与えられる。
以下、Aに対し以下の処理をK回行う。

• 1～Mのいずれかの整数を選び、Aに追加する
• その後、Aの中でX番目に小さい要素を削除する

整数の選び方はM^K通りあるが、全パターンにおいて、最終的にAに残る要素の総和の合計を求めよ。

解法

まずO(NMK)の解法を考える。
i以上の値が、最終的に何個残るかを考える。
dp(i,k,n) := k回処理したとき、i以上の値がn個あるような整数の選び方
状態はO(NMK)通りであり、それぞれの状態からの遷移は、Aに加える値がi未満かi以上かの2択なので、これで全体でO(NMK)で解ける。

これをさらに高速化することを考える。
dp(i,k,n)→dp(i,k+1,m)の状態遷移を考えると、mはnよりXに近づく方向にしか遷移しない。
i以上の値がK回中l回追加されるなら

• nの初期値がX以上の場合、m=max(n-(K-l),X)となる
• nの初期値がX未満なら、m=max(n+l,X)となる

この場合、二項係数の計算によりkに関するループをなくし、O(M(N+K))に抑えることができる。

int N,M,K,X;
int A[202020];
const ll mo=998244353;

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

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;
}

ll from[2020];
ll to[2020];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K>>X;
	FOR(i,N) {
		cin>>A[i];
		A[i]--;
	}
	ll ret=0;
	FOR(i,M) {
		int cur=0;
		FOR(j,N) cur+=A[j]<i;
		ZERO(from);
		
		if(cur>=X-1) {
			FOR(j,K+1) {
				(from[max(cur-j,X-1)]+=comb(K,j)*modpow(M-i,j)%mo*modpow(i,K-j))%=mo;
			}
		}
		else {
			FOR(j,K+1) {
				(from[min(cur+j,X-1)]+=comb(K,j)*modpow(M-i,K-j)%mo*modpow(i,j))%=mo;
			}
		}
		
		
		FOR(j,N+1) {
			(ret+=(N-j)*from[j])%=mo;
		}
	}
	cout<<ret%mo<<endl;
}

まとめ

本番割とギリギリで計算量を落とせてよかった。