Codeforces 235B Let's Play Osu!

問題

n種類のコインを投げて、出た面を順番に記録する。
i番目のコインが表の確率はp[i]である。

このとき、連続する○の大きさの二乗の総和
（たとえば○○×○××だったら、2^2 + 1^2 = 5）の期待値を求めよ。

制約条件

n≦10^5
0≦p≦1

方針

漸化式を立てて計算できないか考える。
i番目までコインを投げたときの列、i=2だとしたら
○○
○×
×○
××
がありえて、それぞれの後ろに○または×が付きうるので、
期待値の計算の一部が再利用できそう。

具体的には期待値から、最後のブロックの寄与分を引いて新たな寄与分を足せばよい。
最後のブロックの寄与分は
i枚コインを投げたとき、最後のブロックの大きさがjである確率をP[i, j]とすれば、
ΣP[i, j] * j^2
これは、ΣP[i, j] * jを計算して保持しておくことで簡単に計算できる。

すなわち、ΣP[i+1, j] * j^2 = ΣP[i, j] * (j+1)^2 * （次のコインが表の確率）
また、ΣP[i+1, j] * j = ΣP[i, j] * (j+1) * （次のコインが表の確率）
ここで、ΣP[i, j] = 1である。

これを使って期待値の漸化式を立てる。
まず、最後に×がくる場合による増減が、Ei * (1 - p[i])
次に新たにできるブロックの寄与分による増減は、ΣP[i+1, j]j^2

Eiの期待値のうち、Ei+1に残る部分は、
後ろのブロックを削り取ったあと、更に○が出たときだけ残るので、
(Ei - ΣP[i, j] * j^2) * p[i]

これらを足し合わせればよい。
editorialはもっときれいなことをやっているぽい。

ソースコード

int n;
double p[100010], a, b, c, d, ans;

int main(){
	cin >> n;
	rep(i, n) cin >> p[i];
	rep(i, n){
		d = (b + 1) * p[i];
		c = (a + 2 * b + 1) * p[i];
		
		ans = ans * (1 - p[i]) + c + (ans - a) * p[i];
		a = c; b = d;
	}
	printf("%.9f\n", ans);
	
	return 0;
}