使用另一个数组中的起始索引对numpy数组的行进行求和

Question

我有一个名为data的NxM numpy数组。 我还有一个长度为N的数组start_indices。 我想要一个新的长度为M的数组，其中第i个元素是sum(data[i][start_indices[i]:])。

以下是一种方法：

import numpy as np
data = np.linspace(0, 11, 12).reshape((3, 4))
data
array([[0, 1, 2, 3],
       [4, 5, 6, 7],
       [8, 9, 10, 11]])
start_indices = np.array([0, 1, 2])
sums = []
for start_index, row in zip(start_indices, data):
    sums.append(np.sum(row[start_index:]))
sums = np.array(sums)

有更numpy风格的方法吗？

- DanielSank

---

你的循环可以重构为列表推导式：np.array([np.sum(row[s:]) for s, row in zip(start,data)])。但速度上没有实质性的差别。 - hpaulj

确实如此..但是zip操作并不是真正需要的，可以直接索引数组而不必创建新数组。 - dermen

Answer 1

您可以创建一个掩码数组。

>>> mask = start_indices[:,None] <= np.arange(data.shape[1])
>>> (data * mask).sum(axis=1)
array([  6.,  18.,  21.])

对于最后一步，您也可以使用np.einsum：

>>> np.einsum('ij,ij->i', data, mask)
array([  6.,  18.,  21.])

尽管在这里使用掩码数组可能效率低下，且迭代过多索引。或者，可以使用 np.fromiter：

>>> it = (r[i:].sum() for r, i in zip(data, start_indices))
>>> np.fromiter(it, data.dtype)
array([  6.,  18.,  21.])

- behzad.nouri

---

在小规模的测试中，所有zip迭代的变体所花费的时间大致相同；这些迭代方法只是语法糖而已。但是带有掩码的einsum比快两倍。我想知道它的扩展性如何。 - hpaulj

这种掩码操作很可爱，但是随着data的大小增加，矩阵乘法会变得非常缓慢。即使是einsum在size(data)增加时也不如列表推导快。我进行了测试，似乎np.fromiter和简单的不使用zip函数的列表推导（我的答案！）是提出的方法中最快的。 - dermen

接受使用fromiter，因为它易于理解并且与hpaulj的cumsum方法一样快（假设我正确使用了timeit）。 - DanielSank

Answer 2

除了zip迭代（多种形式）和掩码求和之外，可能值得尝试的是cumsum。

data[:,::-1].cumsum(axis=1)[range(data.shape[0]), data.shape[1]-1-start_indices]

cumsum 在正确的轴上很容易; 大部分表达式用于提取所需的总和。

在这个小例子中，它比 zip 迭代更快，但比掩码总和慢。但是随着大小的增加，排名可能会改变。

我不认为这些替代方案更“Pythonic”。它们也使用了经过批准的 Python 方法。避免 zip 迭代的方法可能会获得 numpy 的额外加分，但仅当它们在重要的地方提高速度时。

np.reduceat 承诺更好的速度（第一次尝试，不是通用的）：

np.add.reduceat(data.ravel(),[0,4,5,8,10])[::2]

这是一个测试表达式，不考虑生成 indices 列表所需的时间。

indices = np.array([0,4,4,8,8]); indices[::2] += start_indices

- hpaulj

---

reduceat 这个想法很有趣。 cumsum 这个想法几乎是完美的。 - DanielSank

Answer 3

sums = np.array( [data[i, start_indices[i]:].sum() for i in range(data.shape[0])] )

- dermen