具有不同切片的二维 numpy 数组的平均值

Posted 2023-02-23

【中文标题】具有不同切片的二维 numpy 数组的平均值

【英文标题】：Mean over 2d numpy array with varying slices

【发布时间】：2016-10-15 19:51:36

【问题描述】：

我需要计算 2D 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 + 高效）方法是什么？我的数组大小是 ~(50, 50K)。

【问题讨论】：

我选择了最快的时间回答。为了比较和后代： Divakar：12.4 ms Oliver：47.5 ms piRSquared：18.4 ms

【参考方案1】：

你可以使用NumPy broadcasting -

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

【讨论】：

您可以通过将mask.sum(0) 替换为简单的bot_ix 来更优化一步。但迄今为止最好的解决方案！ +1

@OliverW。啊当然！我错过了，谢谢！ OP 应该记下它。

【参考方案2】：

一种方法是让 numpy 计算累积和，然后在新生成的数组中使用花哨的索引，如下所示：

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

我不会对速度做出任何假设，因为它不必要地计算了比严格要求更多的元素的累积和，但这完全取决于您的特定数据。

【参考方案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])