python - 多个切片的 NumPy 平均值

Question

假设我有一个数组a，我想计算其中定义的多个切片的平均值idx：

a = np.arange(10)
idx = np.random.choice([0,1], a.size).astype(bool)

a, idx
Out[1]: (array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
         array([False, False,  True,  True, False, False, False,  True,  True,
                True]))

使用所需的输出：

array([2.5, 8. ])

当然，我可以编写一个简单的for循环，但我更喜欢完全矢量化的方法，因为数组的大小可能会变得非常大并且是二维的。

标签： pythonarraysnumpyslice

Answer 1

可以完全矢量化：

edges = np.diff(idx.astype(np.int8), prepend=0, append=0)
rising = np.where(edges == 1)[0]
falling = np.where(edges == -1)[0]
cum = np.insert(np.cumsum(a), 0, 0)
means = (cum[falling] - cum[rising]) / (falling - rising)

这在我的机器上大约需要 0.2 秒a = np.arange(10**7)。