kmjp's blog

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

TopCoder SRM 786 : Div1 Medium ModCounters

うーん、これはいろいろと思いつかなかった。
https://community.topcoder.com/stat?c=problem_statement&pm=16134&rd=17989

問題

N要素の数列が与えられる。
これらに対し、「ランダムに1要素選んで値をインクリメントする、ただし値512に到達したら0に巻き戻す」という処理をK回行う。
完了後の値の総和の期待値を求めよ。

解法

f(i,c) := i回処理した後、値がcの要素の個数の期待値
とする。f(K,c)さえ求められれば、解は \displaystyle \sum_{c=0}^{511} (c \times f(K,c))となる。

f(i,c)の状態で、値がcの要素が選ばられる確率はf(i,c)/Nであることを考えると、以下の状態遷移が成り立つ。

  • f(i+1,(c+1)%512) += f(i,c)/N
  • f(i+1,c) += f(i,c) - f(i,c)/N

あとは行列累乗の要領で解けるが、512*512の行列で累乗すると微妙に間に合わない。
各行についてみてみると、係数は列がずれてる以外同じなので、1行分についてだけ遷移を求めてみれば十分である。

vector<ll> A;
const ll mo=1000000007;

vector<ll> mulvec(vector<ll>& a,vector<ll>& b) {
	vector<ll> c(a.size(),0);
	int x,y;
	FOR(x,a.size()) FOR(y,b.size()) (c[(x+y)%a.size()]+=a[x]*b[y])%=mo;
	return c;
}

vector<ll> powvec(ll p,vector<ll> a) {
	int i,x,y; vector<ll> r;
	r.resize(a.size());
	r[0]=1;
	while(p) {
		if(p%2) r=mulvec(r,a);
		a=mulvec(a,a);
		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;
}


class ModCounters {
	public:
	int findExpectedSum(vector <int> P, int A0, int X, int Y, int N, int K) {
		A.clear();
		A.push_back(A0);
		while(A.size()<N) A.push_back((A.back()*X+Y)%1812447359);
		while(P.size()<N) P.push_back(A[P.size()]);
		int i,j,x,y;
		
		
		int C[512]={};
		FORR(a,P) C[a%512]++;
		
		vector<ll> V(512);
		V[0]=(1-modpow(N)+mo)%mo;
		V[1]=modpow(N)%mo;
		V=powvec(K,V);
		
		ll ret=0;
		FOR(y,512) FOR(x,512) {
			ret+=V[x]*C[(y-x+512)%512]%mo*y%mo;
		}
		
		
		return ret%mo;
	}
}

まとめ

まずf(i,c)を求めればいいじゃんに到達せず、組み合わせ数でどうにかしようとしてしまって初手が誤り。