yukicoder : No.2571 列辞書順列

問題文の理解にちょっと手間取った。
https://yukicoder.me/problems/no/2571

問題

1～Mの値を取るN要素の整数列Aが与えられる。
以下の2種類のクエリを処理せよ。

• Aの1要素を指定したあたりに上書きする。
• Aの部分列A[L,R]を指定するクエリが与えられるので、長さK以下の、1～Mの値を取る整数列すべてを考えたとき、A[L,R]は何番目か答える。

解法

B=A[L,R]とすると、B[0...(n-1)]=X[0...(n-1)]かつX[n]＜B[n]またはXの長さはnであってX[n]は無いような物は、B[n]*(M^(L-n)-1)/(M-1)である。
これはB[n]*(M^(L-n))の総和とB[n]の総和を取るBITを使えば容易に計算できる。

int N,M,K,Q;
ll A[202020];
ll P[202020];

const ll mo=998244353;
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;
}

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

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K>>Q;
	P[0]=1;
	FOR(i,201010) P[i+1]=P[i]*M%mo;
	FOR(i,N) {
		cin>>A[i];
		A[i]--;
		bt.add(i+1,A[i]*P[N-1-i]%mo);
		bt2.add(i+1,A[i]);
	}
	while(Q--) {
		cin>>i;
		if(i==1) {
			int L,R;
			cin>>L>>R;
			
			if(M==1) {
				cout<<R-L+1<<endl;
			}
			else {
				ll a=mo+bt(R)-bt(L-1);
				a=a*modpow(M,K-N+L)%mo;
				a+=mo-(bt2(R)-bt2(L-1));
				a=a%mo*modpow(M-1)%mo;
				a+=(R-L+1);
				cout<<a%mo<<endl;
			}
		}
		else {
			cin>>x>>y;
			x--;
			bt.add(x+1,mo-A[x]*P[N-1-x]%mo);
			bt2.add(x+1,mo-A[x]);
			A[x]=y-1;
			bt.add(x+1,A[x]*P[N-1-x]%mo);
			bt2.add(x+1,A[x]);
		}
	}
}

まとめ

思ったよりコード短かった。