在python或R中，有没有更有效的方法来枚举离散随机变量的每个可能结果的概率？

Question

我用Python理论计算pmf。这是密码。你知道吗

>>> a_coin = np.array([0,1])
>>> three_coins = np.array(np.meshgrid(a_coin,a_coin,a_coin)).T.reshape(-1,3)
>>> heads = np.sum(three_coins, axis = 1)
>>> df = pd.DataFrame({'heads': heads, 'prob': 1/8})
>>> np.array(df.groupby('heads').sum()['prob'])
array([0.125, 0.375, 0.375, 0.125])

这段代码模拟的是3枚硬币的一次投掷。 可能的结果是{0,1,2,3}。 最后一行代码分别计算每个可能结果的概率。你知道吗

我得放10个硬币进去np.网格（一枚硬币，…，一枚硬币）如果我想计算10枚普通硬币的pmf，这看起来既无聊又低效。你知道吗

问题是，在python或R中有没有更有效的方法来实现这一点？你知道吗

Answer 1

以下是基于fft的numpy解决方案：

import numpy as np
from scipy import fftpack

def toss(n=10, p=0.5):
    t1 = np.zeros(fftpack.next_fast_len(n+1))
    t1[:2] = 1-p, p
    f1 = fftpack.rfft(t1)
    c1 = f1[1:(len(t1) - 1) // 2 * 2 + 1].view(f'c{2*t1.itemsize}')
    c1 **= n
    f1[::(len(t1) + 1) // 2 * 2 - 1] **= n
    return fftpack.irfft(f1)[:n+1]

例如：

>>> toss(3)
array([0.125, 0.375, 0.375, 0.125])
>>> toss(10)
array([0.00097656, 0.00976562, 0.04394531, 0.1171875 , 0.20507813,
       0.24609375, 0.20507813, 0.1171875 , 0.04394531, 0.00976562,
       0.00097656])