AtCoder ABC #215 : G - Colorful Candies 2

力業で通した感はある。
https://atcoder.jp/contests/abc215/tasks/abc215_g

問題

N個の飴があり、それぞれの色が与えられる。
ここからK個選んだ時、飴の色数の期待値を求めよ。

解法

多項式P(x)を、x^nの係数は、飴をn個選んだ時の(色数×組み合わせ)の総和として定める。
そうすると解は[x^K]P(x) / Binom(N,K)となる。

あとはP(x)がどうなるかを考える。
もし、同じ色の飴がm個あるような色を考える。
この飴を1個以上選ぶ場合、P(x)の各係数に寄与することになる。
この飴を1個以上選ぶ場合と、他の飴を0個以上選ぶことを考えると
P(x) += ((1+x)^m-1)*(1+x)^(N-m) = (1+x)^N - (1+x)^(N-m)
となる。この計算は二項係数を前処理しておけばO(N)かかる。

異なる色であっても、mが同じ場合上記加算式は同じ値を加算するだけなので、まとめて行うことができる。
そこで実質mは√N通り程度考えればよく、全体でO(N√N)。

int N;
map<int,int> M;

ll ret[505050];
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;
}

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factmo[NUM_+1];
	if (fact[0]==0) {
		fact[0]=1;
		for (int i=1;i<=NUM_;++i) {
			int x=i;
			factmo[i]=factmo[i-1];
			while(x%mo==0) factmo[i]++, x/=mo;
			fact[i]=fact[i-1]*x%mo;
		}
	}
	if(C_<0 || C_>N_ || factmo[N_]>factmo[C_]+factmo[N_-C_]) return 0;
	return fact[N_]*modpow(fact[C_]*fact[N_-C_])%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>x;
		M[x]++;
	}
	map<int,int> V;
	FORR2(a,b,M) V[b]++;
	
	FOR(i,N+1) ret[i]=M.size()*comb(N,i)%mo;
	FORR2(b,c,V) {
		FOR(i,(N-b)+1) (ret[i]-=c*comb(N-b,i))%=mo;
	}
	
	FOR(i,N+1) if(i) {
		ret[i]=(ret[i]%mo+mo)%mo*modpow(comb(N,i))%mo;
		cout<<ret[i]<<endl;
	}
}

まとめ

O(N√N)だけど、最初3.7sかかってびっくりした。