生成函数

Published: September 06, 2019 by Daltao

Categories:
algorithm 101

概率生成函数

对于取值为自然数的随机变量X，定义概率生成函数为

$f(z)=\sum_{i=0}^\infty P(X=i)z^i$

很显然

$f(1)=\sum_{i=0}^\infty P(X=i)=1$

其中 $P(x=i)$ 为随机变量取到i的概率。利用概率生成函数可以简单地计算期望：

$f'(1)=\sum_{i=1}^\infty P(X=i)iz^{i-1}|_{z=1}=\sum_{i=1}^\infty P(X=i)i=E[X]$

可以递推推出对于 $i\geq 2$ 有

$E[X^i]=(zf^{(i-1)}(z))'|_{z=1}=f^{(i-1)}(z)+zf^{(i)}(z)|_{z=1}=f^{(i-1)}(1)+f^{(i)}(1)$

比较重要的就是计算方差的公式：

$var(X)=E[X^2]-E[X]^2=f''(1)+f'(1)-f'(1)^2$

普通生成函数（ODF）

对于序列 $a_0,a_1,\ldots$ ，它的生成函数为 $f(x)=\sum_{i\geq 0}a_ix^i$ ，记 $[f(x)]_i$ 表示其 $x^i$ 项的系数 $a_i$ 。

两个序列< $a_i$ >与< $b_i$ >的生成函数的加法运算的结果得到的是< $a_i+b_i$ >的生成函数。
两个序列< $a_i$ >与< $b_i$ >的生成函数的乘法运算的结果得到的是< $\sum_{0\leq j\leq i}a_jb_{i-j}$ >的生成函数。

一般我们在做生成函数的运算的时候，采用封闭形式能够加快计算的效率。下面是几个比较常用的封闭形式：

< $1$ >的生成函数为 $\sum_{i\geq 0}x^i=\frac{1}{1-x}$
< $[k|i]$ >的生成函数为 $\sum_{i\geq 0}x^{ki}=\frac{1}{1-x^k}$
< $i$ >的生成函数为 $\sum_{i\geq 0}ix^i=\frac{x}{(1-x)^2}$
< ${n \choose i}$ >的生成函数为 $\sum_{i\geq 0}{n \choose i}x^i=(1+x)^n$
< ${n+i \choose i}$ >的生成函数为 $\sum_{i\geq 0}{n+i \choose i}x^i=\frac{1}{(1-x)^{n+1}}$
< $\frac{1}{i!}$ >的生成函数为 $\sum_{i\geq 0}\frac{1}{i!}x^i=e^x$
< $\frac{1}{i}$ >的生成函数为 $\sum_{i\geq 1}\frac{1}{i}x^i=-\ln(1-x)=\ln(\frac{1}{1-x})$

特别注意由于 $e^x=\sum_{i=0}\frac{x^i}{i!}$ 。因此有 $\sum_{i=k\pmod 2}\frac{x^i}{i!}=\frac{e^x+(-1)^ke^{-x}}{2}$ ，其中 $k=0,1$ 。

题目1：计算 $a_1+\ldots+a_n+b_1+\ldots+b_m+c_1+\ldots+c_k=T$ 的方案数，其中所有变量为非负整数，且 $a_1,\ldots,a_n$ 只能取奇数， $b_1,\ldots,b_m$ 只能取偶数， $c_1,\ldots,c_k$ 为任意非负整数。问有多少种不同的赋值方案，其中 $1\leq n,m,k\leq 100$ ，将结果模上素数 $p$ 。

我们要计算的是生成函数 $f(x)=(\frac{x}{(1-x^2)})^n(\frac{1}{1-x^2})^m(\frac{1}{1-x})^k=\frac{x^n(1+x)^k}{(1-x^2)^{n+m+k}}$ 的第 $T$ 项系数。或者是 $g(x)=\frac{(1+x)^k}{(1-x^2)^{n+m+k}}$ 的第 $T-n$ 按项系数。

进一步展开可以得到：

$g(x)=[\sum_{i\geq 0}{k\choose i}x^i]\cdot [\sum_{i\geq 0}{n+m+k-1+i\choose i}x^{2i}]$

搞出 $\sum_{i\geq 0}{k\choose i}x^i$ ，枚举 $O(k)$ 个 $\sum_{i\geq 0}{n+m+k-1+i\choose i}x^{2i}$ 的系数就可以算出来了。时间复杂度为 $O(k\log_2p+n+m)$ ，还是相当快的。

一道题目。

题目2：给定 $n$ 个非负未知整数 $a_1,\ldots,a_n$ ，要求 $a_i\leq R_i$ ，且 $a_1+\ldots+a_n=T$ 。问有多少种方案，结果模上素数 $p$ ，这里 $R_i\leq 10^6$ ， $T\leq 10^7$ ， $1\leq n\leq 10$ ， $p\leq 10^6$ 。

这个问题可以用DP解决，时间复杂度为 $O(nT)$ 。

也可以用组合数学+容斥解决，时间复杂度为 $O(2^n+p)$ 。

这里讲一下生成函数的做法。可以发现我们要计算的是 $\prod_{i=1}^n\frac{1-x^{R_i+1}}{1-x}=\frac{\prod_{i=1}^n1-x^{R_i+1}}{(1-x)^n}$ 。容易发现上面的分子最多只有 $2^n$ 项的系数非 $0$ ，我们可以用哈希表来表示，计算的时间复杂度为 $O(n2^n)$ 。我们可以将分母展开，得到：

$\frac{\prod_{i=1}^n1-x^{R_i+1}}{\prod_{i=1}^n1-x}=f(x)\sum_{i\geq 0}{n-1+i\choose i}x^i$

我们可以遍历左边的项求解结果，总的时间复杂度为 $O(n2^n+p)$ 。

指数生成函数（EGF）

对于序列 $a_0,a_1,\ldots$ ，其指数生成函数为 $f(x)=\sum_{i\geq 0}\frac{a_i}{i!}x^i$ 。

两个序列< $a_i$ >与< $b_i$ >的生成函数的加法运算的结果得到的是< $a_i+b_i$ >的生成函数。
两个序列< $a_i$ >与< $b_i$ >的生成函数的乘法运算的结果得到的是< $\sum_{0\leq j\leq i}{i \choose j}a_jb_{i-j}$ >的生成函数。

一些常用的指数生成函数和其封闭形式的转换。

< $1$ >的指数生成函数为 $\sum_{i\geq 0}\frac{1}{i!}x^i=e^x$
< $i!$ >的指数生成函数为 $\sum_{i\geq 0}x^i=\frac{1}{1-x}$
< $(i-1)!$ >的指数生成函数为 $\sum_{i\geq 0}\frac{1}{i}x^i=-\ln(1-x)=\ln(\frac{1}{1-x})$

生成函数的系数区间和

给定一个生成函数，我们要计算某个封闭形式为 $f(x)$ 的生成函数，系数和 $\sum_{i=L}^R[f(x)]_i$ 。

这是一个比较经典的技巧，就是我们将 $f(x)$ 乘上 $\sum_{i=0}^{R-L}x^i=\frac{1-x^{R-L+1}}{1-x}$ ，之后取 $x^R$ 项的系数即可。

生成函数解决three sum问题

给定一个长度为 $n$ 的非负序列 $a_1,\ldots,a_n$ ，令 $T(t)$ 表示有多少三元对 $i<j<k$ ，满足 $a_i+a_j+a_k=t$ 。要求计算 $T(0),T(1),\ldots, T(m)$ ，其中 $1\leq n,m\leq 10^5$ 。

这个有很显然的 $O(n^2)$ 暴力做法。但是这里会超时。

我们可以记 $c(i)$ 表示 $i$ 这个数在序列中的出现次数。记 $f(x)=\sum_{i=0}c(i)x^i$ 。考虑 $f^3(x)$ 的含义， $f^3_t$ 表示的是所有的三元对 $i,j,k$ 的数目，满足 $a_i+a_j+a_k=t$ 。和我们的结果非常相近了，但是存在重复统计。我们可以通过容斥统计出所有的三元对 $i,j,k$ 的数目，满足 $a_i+a_j+a_k=t$ 且 $i,j,k$ 互不相同，记它的生成函数为 $g$ 。之后可以发现每个三元对 $i<j<k$ ，都有 $3!$ 种排列，即在 $g$ 中被统计 $3!$ 次，因此我们有 $T(i)=g_i/3!$ 。

生成函数解决乘积和问题

问题1：给定一个长度为 $n$ 的非负序列 $a_1,\ldots,a_n$ ，以及一个素数 $p$ 。要求计算 $\sum_{i=1}^n\sum_{i<j}^n(a_ia_j\pmod p)$ ，注意结果不模 $p$ 。其中 $1\leq n,p\leq 5\times 10^5$

一道原题。

我们不好直接计算，可以发现问题很类似于卷积 $B=A^2$ ，但是卷积的结果为 $B_k=\sum_{i+j=k}A_iA_j$ 。假如我们记 $A_i$ 表示 $i$ 这个数在 $a$ 序列中出现的次数，那么 $B$ 的结果非常接近我们要的结果，当然我们要的实际上是 $C_k=\sum_{ij=k\pmod p}A_iA_j$ 。

需求不满足，我们可以用数学的方式让它满足。由于从 $a$ 中删除所有 $0$ 是不会影响结果的，因此我们可以认为 $a$ 中没有 $0$ 。由于乘法在对数的情况下对应的是加法，即 $\log ab=\log a+\log b$ 。即我们可以找到 $p$ 的任意一个原根 $r$ ，之后令 $A_i$ 表示 $r^i$ 在序列 $a$ 中的出现次数。于是乎在这种意义下 $B=A^2$ ， $B_k$ 表示的是存在多少下标 $i,j$ ，满足 $a_ia_j=r^k\pmod p$ 。而结果为 $\sum_{i=0}^{2p-4}B_i\cdot r^k$ 。但是这里我们下标是可能出现 $i=j$ 的，这一部分可以手动减去即可。

时间复杂度为 $O(p\log_2p+n)$ 。

问题2：给定一个长度为 $n$ 的非负序列 $a_1,\ldots,a_n$ ，以及一个素数 $p$ 。要求计算 $\sum_{i=1}^n\sum_{i<j}^n\lfloor \frac{a_ia_j}{p} \rfloor$ ，注意结果不模 $p$ 。其中 $1\leq n,p\leq 5\times 10^5$

我们考虑问题1中公式的另外一种转换：

$\begin{aligned} \sum_{i=1}^n\sum_{i<j}^n(a_ia_j\pmod p)&=\sum_{i=1}^n\sum_{i<j}^n(a_ia_j-p\lfloor \frac{a_ia_j}{p} \rfloor)\\ &=\sum_{i=1}^n\sum_{i<j}^na_ia_j-p\sum_{i=1}^n\sum_{i<j}^n\lfloor \frac{a_ia_j}{p} \rfloor \end{aligned}\\$

其中

$\sum_{i=1}^n\sum_{i<j}^n\lfloor \frac{a_ia_j}{p} \rfloor=\frac{1}{p}(\sum_{i=1}^n\sum_{i<j}^na_ia_j-\sum_{i=1}^n\sum_{i<j}^n(a_ia_j\pmod p))$

右边公式的两项，前者是非常好算的，维护一个前缀和即可，后者可以通过问题1的方式求解。