我想知道如何使以下代码更短和/或更高效。我可以(或者我应该)通过使用函数方法摆脱for循环,还是我应该使用numpy的方法?
代码计算整数数组的期望值。
vals = np.arange(self.n+1)
# array of probability of each value in vals
parr = np.ones(len(vals))
for i in range(len(vals)):
parr[i] *= self.prob(vals[i])
return np.dot(vals,parr)
根据评论中的要求,方法prob()的实现:
def prob(self, x):
"""Computes probability of removing x items
:param x: number of items to remove
:returns: probability of removing x items
"""
# p is the probability of removing an item
# sl.choose computes n choose x
return sl.choose(self.n, x) * (self.p**x) * \
(1-self.p)**(self.n-x)
答案 0 :(得分:3)
我认为它会更快:
vals = np.arange(self.n+1)
# array of probability of each value in vals
parr = self.prob(vals)
return np.dot(vals,parr)
和功能:
def prob(list_of_x):
"""Computes probability of removing x items
:param list_of_x: numbers of items to remove
:returns: probability of removing x items
"""
# p is the probability of removing an item
# sl.choose computes n choose x
return np.asarray([sl.choose(self.n, e) for e in list_of_x]) * (self.p ** list_of_x) * \
(1-self.p)**(self.n - list_of_x)
因为numpy更快:
import timeit
import numpy as np
list_a = [1, 2, 3] * 1000
list_b = [4, 5, 6] * 1000
np_list_a = np.asarray(list_a)
np_list_b = np.asarray(list_b)
print(timeit.timeit('[a * b for a, b in zip(list_a, list_b)]', 'from __main__ import list_a, list_b', number=1000))
print(timeit.timeit('np_list_a * np_list_b', 'from __main__ import np_list_a, np_list_b', number=1000))
结果:
0.19378583212707723
0.004333830584755033
答案 1 :(得分:2)
循环可以简化为列表理解:
vals = np.arange(self.n+1)
# array of probability of each value in vals
parr = [self.prob(v) for v in vals]
return np.dot(vals, parr)