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より辞書順で小さい値を取る数列を考えると、
となる。
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; }
まとめ
入力がシンプルな問題はテストしやすくて良いね。