使用整数数组作为依赖于列的结束索引的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，则可以使用nansum，但在切片方面，这与等价问题相当。

- Eli S

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

- Austin

Answer 2

一种可能的方法是使用掩码数组：

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]

请注意，当没有好的值时，你会得到一个"missing"值作为结果，这可能对你有用也可能没有用。此外，掩码数组还允许你轻松地进行其他仅适用于有效值(均值等)的操作。

另一个简单的选项是只需将其乘以掩码即可：

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

当没有好的值时，你只会得到零。

- jdehesa

---

还要向Austin的答案喊话，他使用了相同的arange技巧来处理nan值。 - Eli S