AtCoder ARC #124 : E - Pass to Next

これ系のテク、近いものには触れてたのに思出せなかった。
https://atcoder.jp/contests/arc124/tasks/arc124_e

問題

N要素の整数列Aが与えられる。
以下のようにして得られる整数列Bを考える。

• 最初、i番目の人はA[i]個のボールを持つ。
• 各人、手持ちのボールの個数の範囲内で、任意の個数を、同時にi+1番目の人に渡す(N番の人は1番の人に渡す)
• その後、各人が持っているボールの数からなる数列をBとする。

あり得るB全通りのうち、要素の積の総和を求めよ。

解法

Bの組み合わせは、あくまで異なる数列の組み合わせであり、異なる渡し方の組み合わせではない点に注意。
渡し方は異なるが、数列が一致するものを重複して数えない方法として、最低1人はボールを1個も渡さないようにする条件を考える。
この条件を満たす渡し方では、常にBが異なる。

あとは、積の総和の数え方である。
これは以下のように考えることができる。
各人がボールの受け渡しをした後、持っているボールのうち1つを選ぶ選び方、と考えると、この組み合わせはBの積に等しい。
各人が選ぶボールは、もともと自分が持っていたボールか、隣から受け取ったボールのいずれかである。

そこで、以下を状態としてDPを行えばよい。
人が円形に並んでボールを受け渡すことを考えると、以下dp(N, true, true, true) + dp(N, false, false, true)が解。

• dp(n, b, c, z) := 以下のボールの選びかたの組み合わせ
  • n人目までのボールの選択の仕方のうち、
  • b: 先頭の人が選んだボールは、末尾の人から渡されたボールか否か
  • c: n+1人目が選ぶボールは、n人目から渡されたボールか否か
  • z: ここまで渡したボールが0個の人がいたかどうか

渡すボールの数と選び方は、二項係数を使いまとめて数え上げよう。

int N;
ll A[101010];
const ll mo=998244353;
ll dp[101010][2][2];

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 C2(ll a) {
	return a*(a-1)%mo*modpow(2)%mo;
}
ll C3(ll a) {
	return a*(a-1)%mo*(a-2)%mo*modpow(6)%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	ll ret=0;
	FOR(x,2) {
		ZERO(dp);
		dp[0][x][0]=1;
		FOR(i,N) {
			ll P=A[i];
			
			FOR(k,2) {
				// move 0
				(dp[i+1][0][1]+=P*dp[i][0][k])%=mo;
				(dp[i+1][0][1]+=dp[i][1][k])%=mo;
				// move > 0
				(dp[i+1][0][k]+=C2(P)*dp[i][0][k])%=mo;
				(dp[i+1][0][k]+=P*dp[i][1][k])%=mo;
				(dp[i+1][1][k]+=C3(P+1)*dp[i][0][k])%=mo;
				(dp[i+1][1][k]+=(P+C2(P))*dp[i][1][k])%=mo;
			}
		}	
		ret+=dp[N][x][1];
	}
	
	
	cout<<ret%mo<<endl;
	
}

まとめ

ここで覚えておくべきは、数列の積は、数列内で各要素の選択方法の積と言い換えて数え上げられることだな。
一時数え上げではなく、期待値計算で似たような考えの問題も２・３解いていたのに思いつかなかったのが残念。