使用整数数组中与列相关的结束索引的numpy切片

Question

如果我有一个数组，我应用求和

arr = np.array([[1.,1.,2.],[2.,3.,4.],[4.,5.,6]])
np.sum(arr,axis=1)

我得到三行的总数（[4,9,15.]）

我的复杂之处在于，arr包含的数据在某个列索引之后可能不好。我有一个整数数组，它告诉我每行有多少个“好”值，我想对好值求和/平均。说：

ngoodcols=np.array([0,1,2])
np.sum(arr[:,0:ngoodcols],axis=1)  # not legit but this is the idea

如何在循环中实现这一点是很清楚的，但是有没有一种方法可以只求那个数的和，从而产生[0,2,9]，而不用循环呢？等价地，如果我知道如何将列索引中高于b的元素设置为np.nan公司但就切片而言，这几乎是一个等价的问题。你知道吗

Answer 1

这里有一种使用布尔索引的方法。这会将列索引中高于ngoodcols的元素设置为np.nan，并使用np.nansum：

import numpy as np

arr = np.array([[1.,1.,2.],[2.,3.,4.],[4.,5.,6]])
ngoodcols = np.array([0,1,2])

arr[np.asarray(ngoodcols)[:,None] <= np.arange(arr.shape[1])] = np.nan

print(np.nansum(arr, axis=1))
# [ 0.  2.  9.]

Answer 2

一种可能性是使用masked arrays：

import numpy as np

arr = np.array([[1., 1., 2.], [2., 3., 4.], [4., 5., 6]])
ngoodcols = np.array([0, 1, 2])
mask = ngoodcols[:, np.newaxis] <= np.arange(arr.shape[1])
arr_masked = np.ma.masked_array(arr, mask)
print(arr_masked)
# [[     ]
#  [2.0    ]
#  [4.0 5.0  ]]
print(arr_masked.sum(1))
# [  2.0 9.0]

请注意，当这里没有好的值时，您会得到一个“缺少”的值，这可能对您有用，也可能对您不有用。此外，屏蔽数组还允许您轻松地执行仅适用于有效值（平均值等）的其他操作。你知道吗

另一个简单的方法是乘以遮罩：

import numpy as np

arr = np.array([[1., 1., 2.], [2., 3., 4.], [4., 5., 6]])
ngoodcols = np.array([0, 1, 2])
mask = ngoodcols[:, np.newaxis] <= np.arange(arr.shape[1])
print((arr * ~mask).sum(1))
# [0. 2. 9.]

在这里，当没有好的值时，你只得到零。你知道吗