python - 有效地将每一行的元素相乘

Question

给定一个大小为 (n, 3) 且 n 约为 1000 的 ndarray，如何快速地将每一行的所有元素相乘？下面的(不优雅的)第二个解决方案运行时间约为 0.3 毫秒，是否可以改进？

# dummy data
n = 999
a = np.random.uniform(low=0, high=10, size=n).reshape(n/3,3)

# two solutions
def prod1(array):
    return [np.prod(row) for row in array]

def prod2(array):
    return [row[0]*row[1]*row[2] for row in array]

# benchmark
start = time.time()
prod1(a)
print time.time() - start
# 0.0015

start = time.time()
prod2(a)
print time.time() - start
# 0.0003

Answer 1

np.prod 接受轴参数:

np.prod(a, axis=1)

使用 axis=1，计算每一行的列乘积。

完整性检查

assert np.array_equal(np.prod(a, axis=1), prod1(a))

性能

17.6 µs ± 146 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

(1000 倍加速)