[ABC326] E - Revenge of "The Salary of AtCoder Inc."

[Q] https://atcoder.jp/contests/abc326/tasks/abc326_e

考察
1. indexごとに考える。そのindexを踏む期待値はどれだけあるか。

ex[0]: 1/N // 1回目で1/Nを引く
ex[1]: 1/N + 1/NN  // 1回目で1/N + A[0]から1/N
ex[2]: 1/N + A[1]/N + A[0]/N   // 1回目で1/Nと、A[1]から1/Nと、A[0]から1/N
ex[3]: 1/N + A[2]/N + A[1]/N + A[0]/N
... 

2, indexNの期待値ex[N]は、
1/N + (A[0] + A[1] + … + A[N-1]) / N
で求まりそう。
3, 総和の部分(A[0] + A[1] + … + A[N-1])を全部やるとO(N^2)で間に合わないので、別の配列で累積和をもっておく。
4, 1/Nをフェルマーの小定理でやると計算量が多くなるので、invNとしてもっておく。

実装

mint exSum[300030]; // 累積和
mint ex[300030]; // indexを踏む期待値

int main() {
  MOD = 998244353;
  cincout();

  ll N;
  cin >> N;
  vector<ll> A(N);
  rep(i, N) cin >> A[i];

  // 逆元を前計算
  mint invN = 1;
  invN /= N;

  // indexごとに探索
  rep(i, N) ex[i] = invN;
  exSum[0] = ex[0];
  for (ll i = 1; i < N; ++i) {
    ex[i] += exSum[i - 1] * invN;
    exSum[i] += exSum[i - 1] + ex[i];
  }
  mint ans = 0;
  rep(i, N){
    ans += ex[i] * (mint)A[i];
  }
  cout << ans << endl;
}

https://atcoder.jp/contests/abc326/submissions/47037277





いいなと思ったら応援しよう!