numpy数组的和行,其中每个和的起始索引来自另一个数组

Question

我有一个NxM叫做numpy的数组data.我也有一个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)

有更多的numpythonic方式吗？

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.])