QUERGAME - Editorial

PROBLEM LINK:

Div-1 Contest
Div-2 Contest
Div-3 Contest
Practice

Author & Editorialist: Samarth Gupta

DIFFICULTY:

MEDIUM

PREREQUISITES:

Expected Value, FFT/NTT

PROBLEM:

Given an array A, 2 players would delete left element(with probability p_1) or right element(with probability (1-p_1)) alternating turns. You need to find expected sum of numbers with players after K turns.

QUICK EXPLANATION:

We can see that removing elements would generate a sequence of N elements. Find the probability that a_i is on the j^{th} position. For Chef, we sum on odd values of j and for Chefina, we sum on even values. The sum can be found using NTT.

EXPLANATION:

Continuing the quick explanation, let’s find the probability that a_i is on j^{th} position. We can take a_i either by removing elements from left or from right. The probability that a_i is on j^{th} position can then be written as:

p_{i, j} = {j-1 \choose i-1}p_1^ip_2^{j-i} + {j-1 \choose n-i}p_1^{j-1-(n-i)}p_2^{n-i+1}

because in starting (j-1) games, either (i-1) elements from left should have been deleted or (n-i) elements from right.

The expected value in the j^{th} game can then be written as:

E_{j} = \sum\limits_{i=1}^{n}p_{i, j}\times a_i

Next we see that Chef can play at only odd values of j and Chefina would play at even values. Therefore, expected value of Chef after K games can be written as:

E_{chef} = \sum\limits_{j=1, j+=2}^{K} E_{j}

Similarly we can write the expected value for Chefina. Now, in the expression of p_{i, j}, we make the substitution, u = n - i + 1, we can then write E as:

E_{chef} = \sum\limits_{j=1,j+=2}^{K}\sum\limits_{i=1}^{j} {j-1 \choose i-1}p_1^ip_2^{j-i} \times a_i+ {j-1 \choose i-1}p_1^{j-i}p_2^{i} \times a_i

Constructing the Polynomials
Consider P(x) = \sum\limits_{i=1}^{n}\frac{p_1^ia_ix^i}{(i-1)!} and Q(x) = \sum\limits_{i=0}^{n}\frac{p_2^ix^i}{i!}. Multiplying P(x) and Q(x) would give us the first term of summation in E_{chef} divided by (j-1)!. Similarly, we can construct polynomial for the second term. Therefore, we can precalculate answer for every K and answer queries in O(1).

Overall Complexity : O(Nlog(N) + Q)

Problem was made Non-Scorable
Well, this problem was found similar to a problem with a lower constraints. But, I got to know from the codechef team that there was a bonus problem to solve the problem with larger constraints, so it was expected that many people might have already solved it and so the problem was moved to the Non-scorable section.

SOLUTIONS:

Setter's & Editorialist Solution
#include<bits/stdc++.h>
#define ll long long
#define ull unsigned ll
#define uint unsigned
#define pii pair<int,int>
#define pll pair<ll,ll>
#define IT iterator
#define PB push_back
#define fi first
#define se second
#define For(i,j,k) for (int i=(int)(j);i<=(int)(k);i++)
#define Rep(i,j,k) for (int i=(int)(j);i>=(int)(k);i--)
#define CLR(a,v) memset(a,v,sizeof(a));
#define CPY(a,b) memcpy(a,b,sizeof(a));
using namespace std;
const int mo=998244353;
const int FFTN=1<<20;
#define poly vector<int>
namespace FFT{
	int w[FFTN+5],W[FFTN+5],R[FFTN+5];
	int power(int x,int y){
		int s=1;
		for (;y;y/=2,x=1ll*x*x%mo)
			if (y&1) s=1ll*s*x%mo;
		return s;
	}
	void FFTinit(){
		W[0]=1;
		W[1]=power(3,(mo-1)/FFTN);
		For(i,2,FFTN) W[i]=1ll*W[i-1]*W[1]%mo;
	}
	int FFTinit(int n){
		int L=1;
		for (;L<=n;L<<=1);
		For(i,0,L-1) R[i]=(R[i>>1]>>1)|((i&1)?(L>>1):0);
		return L;
	}
	int A[FFTN+5],B[FFTN+5];
	ull p[FFTN+5];
	void DFT(int *a,int n){
		For(i,0,n-1) p[R[i]]=a[i];
		for (int d=1;d<n;d<<=1){
			int len=FFTN/(d<<1);
			for (int i=0,j=0;i<d;i++,j+=len) w[i]=W[j];
			for (int i=0;i<n;i+=(d<<1))	
				for (int j=0;j<d;j++){
					int y=p[i+j+d]*w[j]%mo;
					p[i+j+d]=p[i+j]+mo-y;
					p[i+j]+=y;
				}
			if (d==1<<15)
				For(i,0,n-1) p[i]%=mo;
		}
		For(i,0,n-1) a[i]=p[i]%mo;
	}
	void IDFT(int *a,int n){
		For(i,0,n-1) p[R[i]]=a[i];
		for (int d=1;d<n;d<<=1){
			int len=FFTN/(d<<1);
			for (int i=0,j=FFTN;i<d;i++,j-=len) w[i]=W[j];
			for (int i=0;i<n;i+=(d<<1))	
				for (int j=0;j<d;j++){
					int y=p[i+j+d]*w[j]%mo;
					p[i+j+d]=p[i+j]+mo-y;
					p[i+j]+=y;
				}
			if (d==1<<15)
				For(i,0,n-1) p[i]%=mo;
		}
		int val=power(n,mo-2);
		For(i,0,n-1) a[i]=p[i]*val%mo;
	}
	poly Mul(const poly &a,const poly &b){
		int sza=a.size()-1,szb=b.size()-1;
		poly ans(sza+szb+1);
		if (sza<=30||szb<=30){
			For(i,0,sza) For(j,0,szb)
				ans[i+j]=(ans[i+j]+1ll*a[i]*b[j])%mo;
			return ans; 
		}
		int L=FFTinit(sza+szb);
		For(i,0,L-1) A[i]=(i<=sza?a[i]:0);
		For(i,0,L-1) B[i]=(i<=szb?b[i]:0);
		DFT(A,L); DFT(B,L);
		For(i,0,L-1) A[i]=1ll*A[i]*B[i]%mo;
		IDFT(A,L);
		For(i,0,sza+szb) ans[i]=A[i];
		return ans; 
	}
}
using FFT::Mul;
using FFT::power;
#define mxn 500005
int fact[mxn];
int inv[mxn];
 
void pre()
{
    fact[0] = 1;
    int i;
    For(i,1,mxn-1)
        fact[i] = fact[i-1]*1ll*i%mo;
    inv[mxn-1] = power(fact[mxn-1], mo-2);    
    Rep(i, mxn-2, 0)
        inv[i] = inv[i+1]*1ll*(i+1)%mo;
}
int nCr(int n, int r)
{
    if(n < 0 || r < 0 || n < r)
        return 0;
    int ans = fact[n]*1ll*inv[r]%mo;
    ans = ans*1ll*inv[n-r]%mo;
    return ans;    
}
poly compute(int n, poly &arr, int p1, int p2)
{
    poly pxs(n+1), qxs(n+1), pxe(n+1), qxe(n+1);
    int pw1 = p1, pw2 = p2;
    qxs[0] = 1, qxe[0] = 1;
    For(i, 1, n)
    {
        pxs[i] = (pw1*1ll*inv[i-1]%mo)*1ll*arr[i]%mo;
        pxe[i] = (pw2*1ll*inv[i-1]%mo)*1ll*arr[n-i+1]%mo;
        qxs[i] = (pw2*1ll*inv[i])%mo;
        qxe[i] = (pw1*1ll*inv[i])%mo;
        pw1 = pw1*1ll*p1%mo;
        pw2 = pw2*1ll*p2%mo;
    }
    poly pqs = Mul(pxs, qxs);
    poly pqe = Mul(pxe, qxe);
    poly ans(n+1);
    For(i, 1, n)
        ans[i] = ((pqs[i] + pqe[i])*1ll*fact[i-1])%mo;
    return ans;
}
int main() {
	// your code goes here
	ios_base::sync_with_stdio(false);
    cin.tie(0);
	FFT::FFTinit();
	pre();
	int t;
	cin >> t;
	while(t--)
	{
	    int n, q;
	    cin >> n >> q;
	    vector<int> arr(n+1);
	    int i;
	    For(i, 1, n)
	        cin >> arr[i];
	    int x, y;
	    cin >> x >> y;
	    int p1 = x*1ll*power(y, mo-2)%mo;
	    int p2 = (1 - p1 + mo)%mo;
	   // cout << p1 << " " << p2 << '\n';
	    poly ans = compute(n, arr, p1, p2);
	    int chef[n+1] = {0}, chefina[n+1] = {0};
	    For(i, 1, n)
	    {
	        chef[i] = chef[i-1];
	        chefina[i] = chefina[i-1];
	        if(i%2 == 1)
	            chef[i] = (chef[i] + ans[i])%mo;
	        else
	            chefina[i] = (chefina[i] + ans[i])%mo;
	    }
	    while(q--)
	    {
	        int k;
	        cin >> k;
	        cout << chef[k] << " " << chefina[k] << '\n';
	    }
	}
	return 0;
}
1 Like