python - 对二维数组的每一行的不同切片执行求和

Question

我有一个二维数字数组，想对每一行中的不同索引求平均值。说我有

import numpy as np

data = np.arange(16).reshape(4, 4)

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

我有两个列表，为每一行指定第一个(包含)和最后一个(不包含)索引:

start = [1, 0, 1, 2]
end = [2, 1, 3, 4]

然后我想实现这一点:

result = []
for i in range(4):
    result.append(np.sum(data[i, start[i]:end[i]]))

这使

[1, 4, 19, 29]

但是，我使用的数组比这个例子中的要大很多，所以这个方法对我来说太慢了。有什么聪明的方法可以避免这个循环吗？
我的第一个想法是展平阵列。然后，我想，人们需要以某种方式制作一个切片列表并将其并行应用于数组，我不知道该怎么做。
否则，我正在考虑使用 np.apply_along_axis但我认为这只适用于函数？

Answer 1

让我们带着你的疯狂想法奔跑吧。您可以将数组的索引转换为散乱的索引，如下所示:

ind = np.stack((start, end), axis=0)
ind += np.arange(data.shape[0]) * data.shape[1]
ind = ind.ravel(order='F')
if ind[-1] == data.size:
    ind = ind[:-1]

现在你可以解开原始数组，和 add.reduceat在如此定义的段上:

np.add.reduceat(data.ravel(), ind)[::2] / np.subtract(end, start)

TL;博士

def row_mean(data, start, end):
    ind = np.stack((start, end), axis=0)
    ind += np.arange(data.shape[0]) * data.shape[1]
    ind = ind.ravel(order='F')
    if ind[-1] == data.size:
        ind = ind[:-1]
    return np.add.reduceat(data.ravel(), ind)[::2] / np.subtract(end, start)

计时
使用 @Divakar's answer 中显示的完全相同的数组，我们得到以下结果(当然特定于我的机器):

%timeit einsum_mean(data, start, end)
261 ms ± 2.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit broadcasting_mean(data, start, end)
405 ms ± 1.64 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit ragged_mean(data, start, end)
520 ms ± 3.68 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit row_mean(data, start, end)
45.6 ms ± 708 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

有点令人惊讶的是，尽管通过将感兴趣区域之间的所有数字相加做了很多额外的工作，但这种方法的运行速度比所有其他方法快 5-10 倍。原因可能是它的开销极低:索引数组很小，并且只对一维数组进行一次传递。