python - 逐行计算二维 np 数组的平均值，直到某个不同的元素

Question

让我们考虑一个二维数组

arr = np.array([[ 0,  1,  2,  3,  4],
                [ 5,  6,  7,  8,  9],
                [10, 11, 12, 13, 14],
                [15, 16, 17, 18, 19],
                [20, 21, 22, 23, 24]])

让我们也考虑另一个包含索引限制的数组。

limits = np.array([2, 3, 2, 4, 3])

有没有办法计算前一个数组的平均值（按行），但只考虑每行的一堆（由第二个数组给出）元素？特别是，在我之前的示例中，我希望计算均值时仅考虑第一行的前 2 个元素，第二行的前 3 个元素，第三行的前 2 个元素，依此类推.

我提出了一个非常基本的解决方案，使用 for 循环：

arr_means = np.zeros(5, dtype=np.float32)
for i, lim in zip(range(5), limits):
    arr_means[i] = arr[i, :lim].mean()

它返回以下内容：

np.array([ 0.5,  6. , 10.5, 16.5, 21. ], dtype=float32)

是否可以对这段代码进行矢量化？

标签： pythonarraysnumpy

Answer 1

您可以为要平均的位置创建一个掩码：

mask = np.arange(arr.shape[1])<limits[...,None]

out = (arr*mask).sum(1)/mask.sum(1)

输出：

array([ 0.5,  6. , 10.5, 16.5, 21. ])