「洛谷P6667」[清华集训2016]如何优雅地求和 题解

题目

提供一种概率生成函数的做法

首先简单介绍一下概率生成函数：

概率生成函数是一类用于解决、分析在自然数域上的离散随机变量相关问题的普通生成函数，其形式为 \(F(x)=\sum\limits_{i=0}^\infty\operatorname{P}(X=i)x^i\)

有如下公式：

\(\operatorname{E}(X)=F^\prime(1)\)

\(\operatorname{E}(X^{\underline{n}})=F^{(n)}(1)\)

对于成功概率为 \(p\)，重复试验 \(n\) 次的二项分布，概率生成函数为 \((1-p+px)^n\)

下面是这题的做法

思路

发现这个式子很像二项分布，事实上，这道题就是对于成功概率 \(p\)，重复试验 \(n\) 次的二项分布成功次数 \(X\)，求 \(\operatorname{E}[f(X)]\)

根据期望的线性性，可以拆开，又有 \(\operatorname{E}(X^{\underline{k}})=F^{(k)}(1)=p^kn^{\underline{k}}\)（手动求导可得），其中 \(F(x)\) 是该二项分布的概率生成函数

考虑将 \(f(x)\) 转为下降幂多项式 \(g(x)=\sum_{i=0}^mg_ix^{\underline{i}}\)，答案即为 \(\sum_{i=0}^mg_ip^in^{\underline{i}}\)

由于已经给出了 \(0\sim m\) 处的点值，直接将 \(\sum\frac{f(i)}{i!}\) 和 \(\sum\frac{(-1)^i}{i!}\) 卷积的 \(i\) 次项即为 \(g_i\)（不会的可以先看这道模板题）

时间复杂度 \(\operatorname{O}(m\log m)\)

代码

const int N = 20005;
const int M = 100005;
const int Mod = 998244353, g = 3;

class Poly {
public:
  int *a;
  Poly() {
    a = new int[M];
    memset(a, 0, sizeof(int) * M);
  }
  void NTT(int lim, bool opt);
};
int r[M];

int pow(int a, int b, int m);

int pw[N], dw[N], fac[N], ifac[N];
int f[N];
int n, m, p;

int main () {
  read(n), read(m), read(p);
  int lim = 1;
  while (lim <= (m << 1)) {
    lim <<= 1;
  }
  for (int i = 1; i < lim; ++i) {
    r[i] = (r[i >> 1] >> 1) | ((i & 1) * (lim >> 1));
  }
  pw[0] = 1;
  dw[0] = 1;
  fac[0] = 1;
  for (int i = 1; i <= m; ++i) {
    pw[i] = pw[i - 1] * 1ll * p % Mod;
    dw[i] = dw[i - 1] * 1ll * (n - i + 1) % Mod;
    fac[i] = fac[i - 1] * 1ll * i % Mod;
  }
  ifac[m] = pow(fac[m], Mod - 2, Mod);
  for (int i = m - 1; i >= 0; --i) {
    ifac[i] = ifac[i + 1] * 1ll * (i + 1) % Mod;
  }
  for (int i = 0; i <= m; ++i) {
    read(f[i]);
  }
  Poly F, G;
  for (int i = 0; i <= m; ++i) {
    F.a[i] = f[i] * 1ll * ifac[i] % Mod;
    G.a[i] = ((i & 1) ? (Mod - ifac[i]) : ifac[i]);
  }
  F.NTT(lim, false), G.NTT(lim, false);
  for (int i = 0; i < lim; ++i) {
    F.a[i] = F.a[i] * 1ll * G.a[i] % Mod;
  }
  F.NTT(lim, true);
  int ans = 0;
  for (int i = 0; i <= m; ++i) {
    ans = (ans + F.a[i] * 1ll * pw[i] % Mod * dw[i] % Mod) % Mod;
  }
  write(ans), EL;
  return 0;
}

int pow(int a, int b, int m) {
  int ans = 1, now = a;
  while (b) {
    if (b & 1) {
      ans = ans * 1ll * now % m;
    }
    now = now * 1ll * now % m;
    b >>= 1;
  }
  return ans;
}
void Poly::NTT(int lim, bool opt) {
  for (int i = 0; i < lim; ++i) {
    if (i < r[i]) {
      std::swap(a[i], a[r[i]]);
    }
  }
  for (int l = 1; l < lim; l <<= 1) {
    int gn = pow(g, (Mod - 1) / (l << 1), Mod);
    for (int i = 0; i < lim; i += (l << 1)) {
      for (int j = 0, gw = 1; j < l; ++j, gw = gw * 1ll * gn % Mod) {
        int x = a[i + j], y = a[i + j + l] * 1ll * gw % Mod;
        a[i + j] = (x + y) % Mod;
        a[i + j + l] = (x - y + Mod) % Mod;
      }
    }
  }
  if (opt) {
    std::reverse(a + 1, a + lim);
    int iv = pow(lim, Mod - 2, Mod);
    for (int i = 0; i < lim; ++i) {
      a[i] = a[i] * 1ll * iv % Mod;
    }
  }
}