kmjp's blog

競技プログラミング参加記です

yukicoder : No.2068 Restricted Permutation

1ステップずつ丁寧に行えば難しくない?
https://yukicoder.me/problems/no/2068

問題

整数N,K,Xが与えられる。
長さNの順列Pに対し、f(P)とは長さNの順列のうちPより辞書順で手前に来るものを差す。

P[K]=Xである順列Pに対し、f(P)の総和を求めよ。

解法

以下1-originで記載する。
K[i]を、Pを先頭から見て行ったとき、P[i]がP[i]...P[N]の中で小さい順に何番目かを示す数列とする。
この時、f(P)は、i番目の要素で初めてPより辞書順で小さい値を取る数列を考えると、
 \displaystyle f(P) = \sum_i (K_i-1)(N-i)!となる。

Pが、問題の条件を満たす(N-1)!通りの数列を等確率で1つ選ぶとし、各iごとに(K_i-1)(N-i)!の期待値を考える。(最後に数列の個数(N-1)!を掛ければ、個数になる)
iがKより大きいケースでは、K番目まで等しい2つの数列Pがあったとき、片方がもう片方より辞書順で小さい確率は1/2なので、これらの総和はまとめて(N-K)!/2と計算できる。

iがK未満のケースは、P[i]が任意の値を取れるケースからXを取るケースを引くと良い。

int N,K,X;
const ll mo=998244353;

const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

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;

	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;
	
	cin>>N>>K>>X;
	if(N==1) {
		cout<<0<<endl;
		return;
	}
	ll ret=0;
	
	for(i=1;i<=K;i++) {
		if(i==K) {
			ll sub=(X-1-1LL*(i-1)*(X-1)%mo*modpow(N-1)%mo+mo)*fact[N-i]%mo;
			ret+=sub;
		}
		else {
			ll tot=fact[N-i+1]%mo*(mo+1)/2%mo;
			ll sub=(X-1-1LL*(i-1)*(X-1)%mo*modpow(N-1)%mo+mo)*fact[N-i-1]%mo;
			ret+=tot+mo-sub;
		}
	}
	ret+=(fact[N-K]-1)*(mo+1)/2%mo;
	
	
	cout<<ret%mo*fact[N-1]%mo<<endl;
	
}

まとめ

入力がシンプルな問題はテストしやすくて良いね。