AtCoder ABC #231 (パナソニックプログラミングコンテスト2021) : G - Balls in Boxes

これ系は何度か解いたのでまぁ。
https://atcoder.jp/contests/abc231/tasks/abc231_g

問題

N個の箱があり、i番目の箱にはA[i]個のボールが入っている。
ここに、K個の区別できるボールをそれぞれ等確率でいずれかの箱に追加する。

その後の各箱のボール数からなる数列をBとしたとき、Bの積の期待値を求めよ。

解法

Bの積を計算するというのは、各箱から1個ずつボールを選ぶ組み合わせを選ぶのと等しい。
そこで、元からあったボールを選ぶ箱と、追加したボールを選ぶ箱を区別して考える。

f(n,k) := n個目までの箱のうち、追加したボールを選ぶ箱がk個である組み合わせ
とすると、f(N,k)はO(N^2)で埋められる。

次にk=0～Nに対し、Bの積を考えよう。
K個のうちあるk個が、k個の箱にそれぞれ1個ずつ入ればf(N,k)分だけBの積の総和に寄与する。
k個の選び方はP(N,k)通り、その通り箱に入る確率は((1/N)^k)なので、f(N,k)*P(N,k)*((1/N)^k)だけ解に計上すればよい。

int N,K;
int A[1010];
const ll mo=998244353;
ll dp[1010][1010];

ll modpow(ll a, ll n=mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll comb(ll P_,ll Q_) {
	if(P_<0 || Q_<0 || Q_>P_) return 0;
	ll p=1,q=1;
	for(int i=1;i<=Q_;i++) p=p*P_%mo, q=q*i%mo,P_--;
	
	return p;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	
	dp[0][0]=1;
	FOR(i,N) {
		cin>>A[i];
		FOR(j,i+1) {
			// new
			(dp[i+1][j+1]+=dp[i][j])%=mo;
			// exist
			(dp[i+1][j]+=dp[i][j]*A[i])%=mo;
		}
	}
	
	ll ret=0;
	for(i=0;i<=N;i++) {
		ret+=dp[N][i]*comb(K,i)%mo*modpow(modpow(N),i)%mo;
	}
	
	cout<<ret%mo<<endl;
	
}

まとめ

最初追加したボールを区別しない問題を解いてしまった。
こちらも同じ考え方で、大差ない難易度で解くことができる。