在具有不同切片的2D numpy数组上求平均值

Question

我需要计算一个二维numpy数组每列的平均值，其中按列切片会改变。

例如，我有一个数组

    arr = np.arange(20).reshape(4, 5)

每列切片的结束索引定义为平均值

    bot_ix = np.array([3, 2, 2, 1, 2])

第一列的平均值将会是：

    arr[0:bot_ix[0], 0].mean()

什么是适当的（即Pythonic + efficient）方法来做到这一点？我的数组大小约为（50, 50K）。

- Jesse Lopez

---

我选择了最快的答案。供比较和后世参考： Divakar：12.4毫秒 Oliver：47.5毫秒 piRSquared：18.4毫秒 - Jesse Lopez

Answer 1

你可以使用 NumPy广播 -

mask = bot_ix > np.arange(arr.shape[0])[:,None]
out = np.true_divide(np.einsum('ij,ij->j',arr,mask),mask.sum(0))

用于验证结果的示例运行 -

In [431]: arr
Out[431]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19]])

In [432]: bot_ix
Out[432]: array([3, 2, 2, 1, 2])

In [433]: np.true_divide(np.einsum('ij,ij->j',arr,mask),mask.sum(0))
Out[433]: array([ 5. ,  3.5,  4.5,  3. ,  6.5])

In [434]: [arr[0:item, i].mean() for i,item in enumerate(bot_ix)]
Out[434]: [5.0, 3.5, 4.5, 3.0, 6.5] # Loopy version to test out o/p

- Divakar

---

你甚至可以通过用 bot_ix 替换 mask.sum(0) 来进一步优化。但到目前为止这是最好的解决方案！+1 - Oliver W.

@OliverW。当然！我错过了那个，谢谢！OP应该注意一下。 - Divakar

Answer 2

一种方法是让numpy计算累积和，然后在新生成的数组中使用高级索引，像这样：

np.true_divide(arr.cumsum(axis=0)[bot_ix-1,range(arr.shape[1])], bot_ix)

我不会对速度做出任何假设，因为这样会无谓地计算比实际需要更多的元素的累积和，但它完全取决于你特定的数据。

- Oliver W.

Answer 3

Divakar和Oliver W.的结合。

mask = np.arange(arr.shape[0])[:, None] < bot_ix
(arr * mask).sum(0) / bot_ix.astype(float)

array([ 5. ,  3.5,  4.5,  3. ,  6.5])

- piRSquared