kmjp's blog

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

AtCoder ARC #111 : F - Do you like query problems?

あと一歩だったんだけどな。
https://atcoder.jp/contests/arc111/tasks/arc111_f

問題

N要素の数列Aを考える。初期値は全要素0である。また、解となる変数ansの初期値を0にする。
ここで、以下の3通りのクエリが計Q回来ることを考える。

  • 区間[L,R]とm未満の非負整数整数vが与えられるので、A[L...R]のうちv以下のものをvにする
  • 区間[L,R]とm未満の非負整数整数vが与えられるので、A[L...R]のうちv以上のものをvにする
  • 区間[L,R]が与えられるので、A[L...R]をansに加える。

され、クエリ1回はN*(N+1)/2*(2*M+1)通り考えられる。
全クエリで(N*(N+1)/2*(2*M+1))^Qの全組み合わせが行われたとき、ansの合計を求めよ。

解法

各要素毎に、ansへの寄与の期待値を求めよう。
その和を取って(N*(N+1)/2*(2*M+1))^Qを掛ければ解になる。

まず重要な考察として、任意のタイミングである変数の値が1~(m-1)となる確率は等しい。
区間は置いておいて、(2M+1)通りのクエリの結果変数の値がx以外からxになる確率はいずれも(M-1)/(2M+1)で等しい。
そのため、初期値である0以外に到達する確率は常時等しくなる。
そこで、数列の値がある値である確率は、0か0でないかの2値だけ考えればよく、0でないなら1~(m-1)が等確率なので平均してm/2が加算されると考えてしまってよい。

i番目の要素が、区間[L,R]に含まれる確率P(i)=*1/(N(N+1)/2)とする。
Q回のクエリにおいて

  • 変数の値が0→1以上になる確率はP(i)*(M-1)/(2M+1)
  • 変数の値が0→0になる確率は1-P(i)*(M-1)/(2M+1)
  • 変数の値が1以上→0になる確率はP(i)/(2M+1)
  • 変数の値が1以上→1以上になる確率は1-P(i)/(2M+1)
  • 変数の値が1以上で、ansに加算される確率はP(i)/(2M+1)

で遷移する。
そこで、変数の値が0である確率と1以上である確率、あとansの期待値の3変数に対し、状態遷移を3*3の行列で表現すれば行列累乗テクでQ回分のクエリのansへの加算分がO(logQ)で求められる。

iごとにP(i)が異なるので、iを総当たりすれば全体をO(NlogQ)で処理できる。

int N,M,Q;
const ll mo=998244353;

ll dp[202020][2];

const int MAT=3;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo;
	return r;
}

Mat powmat(ll p,Mat a,int n=MAT) {
	int i,x,y; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(i,n) r.v[i][i]=1;
	while(p) {
		if(p%2) r=mulmat(r,a,n);
		a=mulmat(a,a,n);
		p>>=1;
	}
	return r;
}

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;
	
	cin>>N>>M>>Q;
	
	ll P=1LL*N*(N+1)/2%mo;
	ll ret=0;
	
	FOR(i,N) {
		ll T=1LL*(i+1)*(N-i)%mo*modpow(P)%mo;
		ll godown=1*modpow(2*M+1)%mo*T%mo;
		ll goup=(M-1)*modpow(2*M+1)%mo*T%mo;
		Mat A;
		A.v[0][0]=1+mo-goup;
		A.v[0][1]=goup;
		A.v[1][0]=godown;
		A.v[1][1]=1+mo-godown;
		A.v[1][2]=godown;
		A.v[2][2]=1;
		A=powmat(Q,A);
		(ret+=A.v[0][2])%=mo;
	}
	ll C=1LL*M*(mo+1)/2%mo;
	ret=ret*C%mo;
	ret=ret*modpow(P*(2*M+1)%mo,Q)%mo;
	cout<<ret<<endl;
	
	
}

まとめ

O(NQ)解法を作ってから、Qの分ではなくNの分を削減する方向に走ったのが失策。

*1:i+1)*(N-i