AtCoder ARC #174 : E - Existence Counting

なんか運よく単独全完できた回。
https://atcoder.jp/contests/arc174/tasks/arc174_e

問題

整数N,Kが与えられたとき、以下の数列Aを考える。

• |A|=K
• Aの値はdistinctで、1～Nの値を取る。

Aの一種であるPが与えられる。
Aのうち以下を満たすものは何通りか。
t=1～Nそれぞれに対し求めよ。

• A中にtを含む
• Aは辞書順でP以下である。

解法

AとPがi要素目で異なる、すなわち

• A[0...(i-1)]=P[0...(i-1)]
• A[i]＜P[i]

となることを考える。
iごとに、A[i+1]以降の決め方をそれぞれBITで管理することを考える。

• A[0...(i-1)]にtが含まれる場合、A[i+1]以降は何でも良い。
• A[0...(i-1)]にtが含まれない場合、A[i]=tにしたいならP[i]＞tでなければならない。
• A[i+1]以降にtを置くなら、tを置く場所は(K-(i+1))通りの候補がある。

tを動かしながら、上記3値を管理するBITの値を差分更新していこう。

int N,K;
int P[303030],R[303030];
const ll mo=998244353;

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

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s%mo;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<ll,20> bt,did,bt2;



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;
	
	MINUS(R);
	cin>>N>>K;
	FOR(i,N) did.add(i,1);
	FOR(i,K) {
		cin>>P[i];
		P[i]--;
		R[P[i]]=i;
		did.add(P[i],-1);
		x=did(P[i]);
		//初回で登場
		bt.add(i,fact[N-1-i]%mo*factr[N-K]%mo);
		//2回目以降で登場
		if(i<K-1) {
			bt.add(i,1LL*(x-1)*(K-1-i)%mo*fact[N-2-i]%mo*factr[N-K]%mo);
		}
		bt2.add(i,x*fact[N-1-i]%mo*factr[N-K]%mo);
			
	}
	
	FOR(i,N) {
		if(R[i]==-1) {
			cout<<bt(K-1)%mo<<endl;
		}
		else {
			x=R[i];
			bt.add(x,mo-fact[N-1-x]%mo*factr[N-K]%mo);
			if(x<K-1) bt.add(x,(K-1-x)%mo*fact[N-2-x]%mo*factr[N-K]%mo);
			
			ll a=bt(x)+bt2(K)-bt2(x)+mo+1;
			cout<<a%mo<<endl;
		}
	}
	
}

まとめ

この回CDでえらく手間取ったけど、EFはすんなり行けたのでよかった。