查找NumPy数组与值列表中的任何值相等的位置 [英] Find where a NumPy array is equal to any value in a list of values

问题描述

我有一个整数数组，想要找到该数组等于多个值列表中任何值的位置.

I have an array of integers and want to find where that array is equal to any value in a list of multiple values.

可以通过分别处理每个值或在循环中使用多个或"语句来轻松完成此操作，但是我觉得必须有一种更好/更快的方法.我实际上正在处理大小为4000 x 2000的数组，但这是问题的简化版:

This can easily be done by treating each value individually, or by using multiple "or" statements in a loop, but I feel like there must be a better/faster way to do it. I'm actually dealing with arrays of size 4000 x 2000, but here is a simplified edition of the problem:

fake = arange(9).reshape((3,3))

array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

want = (fake==0) + (fake==2) + (fake==6) + (fake==8)

print want 

array([[ True, False,  True],
       [False, False, False],
       [ True, False,  True]], dtype=bool)

我想要的是一种从包含fake和值列表[0, 2, 6, 8]的单个命令中获取want的方法.

What I would like is a way to get want from a single command involving fake and the list of values [0, 2, 6, 8].

我假设有一个已经包含了这个包的软件包，这比我刚用Python编写一个带有循环的函数要快得多.

I'm assuming there is a package that has this included already that would be significantly faster than if I just wrote a function with a loop in Python.

推荐答案

函数 numpy.in1d 似乎可以满足您的要求.唯一的问题是它仅适用于一维数组，因此您应该像这样使用它:

The function numpy.in1d seems to do what you want. The only problems is that it only works on 1d arrays, so you should use it like this:

In [9]: np.in1d(fake, [0,2,6,8]).reshape(fake.shape)
Out[9]: 
array([[ True, False,  True],
       [False, False, False],
       [ True, False,  True]], dtype=bool)

我不知道为什么只限于一维数组.看看它的源代码，它首先看起来像展平两个数组，然后执行一些巧妙的排序技巧.但是没有什么能阻止它再次使结果变得平坦，就像我在这里必须手动做的那样.

I have no clue why this is limited to 1d arrays only. Looking at its source code, it first seems to flatten the two arrays, after which it does some clever sorting tricks. But nothing would stop it from unflattening the result at the end again, like I had to do by hand here.

这篇关于查找NumPy数组与值列表中的任何值相等的位置的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！