计算“多项式系数”的numpy/scipy函数

Question

有没有什么Python函数（可能是来自numpy或scipy的）可以计算在这个公式中(1+x+x**2+x**3+...+x**(k-1))**n展开时，x**r的系数？这里的条件是k>=1，n>=0，并且0<=r<=n(k-1)。

这个系数有时被称为多项式系数（PC）（比如可以参考这里）。

如果没有这样的函数，你能想到一个高效的计算方法吗？（我对简单粗暴的方法不感兴趣）。

Answer 1

你实际上是在对一组全是1的数组进行n次卷积，也就是把这些1叠加起来。
所以，你有没有考虑过使用快速傅里叶变换（FFT）来处理一个足够填充零的数组，把它的元素提升到n次方，然后再用逆FFT来找出所有的系数呢？

(1+x+x**2+x**3+...+x**(k-1))**n

然后你只需要读取你感兴趣的那些系数就可以了。

更新：

因为FFT是循环的，所以你需要一个数组，它的大小至少要和

(1+x+x**2+x**3+...+x**(k-1))**n

中的项数一样，换句话说，就是要有 (k-1)*n+1 的长度，这样结果才不会在末尾绕回去（或者至少在绕回去时只会给受影响的元素加上零）。通常，这个数组的长度最好是2的幂，因为FFT算法需要这样（如果实现不要求这样，它会在你的输入后面加零直到满足条件）。

用类似C语言的伪代码表示：

unsigned int m = 1;
while(m<(k-1)*n+1) m *= 2;

complex c[m];
for(unsigned int i=0;i!=k;++i) c[i] = complex(1.0, 0.0);
for(unsigned int i=k;i!=m;++i) c[i] = complex(0.0, 0.0);

c = fft(c);
for(unsigned int i=0;i!=m;++i) c[i] = pow(c[i], double(n));
c = inv_fft(c);

最后，复数数组 c 中的第 r 个元素的实部就是 x**r 的系数，虚部为零。
由于这一切都是用浮点数来计算的，你需要注意这些元素可能会有累积的舍入误差。你可以通过把它们四舍五入到最近的整数来部分修正这个问题，但要知道对于足够大的 k 和 n，这些误差可能会超过0.5，这样可能会导致结果有一些小的相对误差。

在网上快速搜索一下，你会发现numpy有FFT及其逆变换的实现，分别是 numpy.fft.rfft 和 numpy.fft.irfft，你可以在输入数据为实数时使用它们。