Numpy:多个值的索引

Question

我想通过高效索引将2D数组的一些值设置为特定的数字。

假设我有一个2D numpy数组，

A = array([[1, 6, 6],
[9, 7, 7], 
[10, 2, 2]])

我想得到数组中属于一组数字的索引，比如indList=[10, 1]，这样我就可以将它们设置为零。然而，indList可能是一个巨大的列表。

在没有for循环的情况下，有没有更快的方法可以做到这一点？

作为一个for循环，

indList = [10, 1]
for i in indList:
A[A==i] = 0

但当indList较大时，这可能会变得效率低下。

Answer 1

使用numpy，您可以通过首先查找indList中元素的索引，然后将其设置为零来将其向量化。

A = np.array([[1, 6, 6],
[9, 7, 7],
[10 ,2 ,2]])
A[np.where(np.isin(A, [10,1]))] = 0

这提供

A = [[0 6 6]
[9 7 7]
[0 2 2]]

Answer 2

根据@Miket25的回答，实际上不需要添加np.where层。np.isin(A, [10, 1])返回一个布尔数组，该数组完全可以作为索引。所以简单地做

A[np.isin(A, [10, 1])] = 0