B - Sample Game(dp)
(dp_i)代表在位置(i)时还有多少步结束的期望值。
因为期望具有线性性,根据下一个位置(j),有两种情况:
- (j<i),说明下一步就结束了,有(dp_i=p_i imes 1)
- (j>=i),则说明还有(dp_j)步才能结束,有(dp_i=p_j(dp_j+1))
可得
(dp_i=sumlimits_{j<i}{p_j}+sumlimits_{jge i}{p_j(dp_j+1)})
移项可得
(dp_i=frac{sumlimits_{j<i}{p_j}+sumlimits_{j>i}{p_j(dp_j+1)+p_i}}{1-p_i})
由于(E((x+1)^2)=E(x^2+2x+1)=E(x^2)+2E(x)+1)
可以根据(dp)同时维护第二个数组代表平方的期望。
(i)从大到小,(i=0)处的值即为答案。时间复杂度O(n^2)
用生成函数的方法可以O(n)解决,还不会。
#include <bits/stdc++.h>
#define endl '
'
#define IOS std::ios::sync_with_stdio(0); cin.tie(0); cout.tie(0)
#define mp make_pair
#define seteps(N) fixed << setprecision(N)
typedef long long ll;
using namespace std;
/*-----------------------------------------------------------------*/
ll gcd(ll a, ll b) {return b ? gcd(b, a % b) : a;}
#define INF 0x3f3f3f3f
const int N = 3e5 + 10;
const int M = 998244353;
const double eps = 1e-5;
inline ll qpow(ll a, ll b, ll m) {
ll res = 1;
while(b) {
if(b & 1) res = (res * a) % m;
a = (a * a) % m;
b = b >> 1;
}
return res;
}
ll dp1[N], dp2[N];
ll p[N], pre[N];
int main() {
IOS;
int n;
cin >> n;
ll sum = 0;
for(int i = 1; i <= n; i++) {
cin >> p[i];
sum = (sum + p[i]) % M;
}
ll rsum = qpow(sum, M - 2, M);
for(int i = 1; i <= n; i++) {
p[i] = p[i] * rsum % M;
pre[i] = (pre[i - 1] + p[i]) % M;
}
for(int i = n; i >= 0; i--) {
ll sumdp1 = 0, sumdp2 = 0;
for(int j = i + 1; j <= n; j++) {
sumdp1 = (sumdp1 + p[j] * (dp1[j] + 1) % M) % M;
sumdp2 = (sumdp2 + p[j] * (dp2[j] + 2 * dp1[j] + 1) % M) % M;
}
ll q = qpow((1 - p[i]) % M + M, M - 2, M);
dp1[i] = (sumdp1 + pre[i - 1] + p[i]) * q % M;
dp2[i] = (sumdp2 + pre[i - 1] + p[i] * (1 + 2 * dp1[i]) % M) * q % M;
}
cout << dp2[0] << endl;
}