优化几何分布计算中的琐碎求和

Question

def geo_dist_names(p, k):
    sum = 0
    for i in range(1, 4):
        sum += p**i   
    return (p**(1+k))/sum

p 是 0 到 1 之间的浮点数，k 是 0 到 3 之间的整数。 该函数基本上只是找到与给定 p 和 k 相关的几何分布中的值，然后通过除以 的 4 个潜在值的总和来对其进行归一化k。。

它有效，但我多次调用这个函数，所以我想知道是否有更优化的方法来执行此操作？

Answer 1

你的代码的矢量版本是：

import numpy as np

def geo_dist_names(p, k):
    return (p**(1+k))/(p**np.arange(1,4)).sum()

然而，我不确定它会比纯 python 更快，因为这里的范围很小，所以 numpy 的开销可能不可忽略。

编辑。确实，假设：

def geo_dist_names_python(p, k, N=4):
    sum = 0
    for i in range(1, N):
        sum += p**i   
    return (p**(1+k))/sum

def geo_dist_names_numpy(p, k, N=4):
    return (p**(1+k))/(p**np.arange(1,N)).sum()

numpy 仅在范围增加时更好：

Answer 2

不知道有没有更优化。但作为一句单行话，你可以这样写：

def geo_dist_names(p, k):
    return (p**(1+k))/(p+p**2+p**3)