我有一个NxM名为的numpy数组data。我也有一个N名为的长度数组start_indices。我想要一个新的长度M数组，其中ith元素是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方式吗？
您可以创建遮罩数组
>>> 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.])
本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句