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

时间:2019-03-06 17:03:04

标签: python arrays numpy indexing

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

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,但是就切片而言,这几乎是等效的问题。

2 个答案:

答案 0 :(得分: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.]

答案 1 :(得分:1)

一种可能是使用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.]

在这里,当没有好的价值时,您只会得到零。