生成函数

Published: by Creative Commons Licence

  • Tags:

概率生成函数

对于取值为自然数的随机变量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})$

题目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\mod 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\mod 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\mod 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\mod 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}\\ \Rightarrow \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\mod p))\]

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

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

参考资料