「洛谷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;
    }
  }
}

本博客采用知识共享署名-相同方式共享 4.0 国际许可协议进行许可
本文链接:https://blog.seniorious.cc/2020/luogu-6667/