获取两个数组满足条件的索引

3 投票
4 回答
18007 浏览
提问于 2025-04-16 20:47

我有两个numpy数组:ab。我想找出所有满足条件的索引,也就是那些a == 1并且b == 0的地方。

比如说,如果我有下面这样的数组:

a = [0, 1, 3, 5, 1, 1, 2]

b = [1, 0, 2, 5, 3, 0, 6]

我希望能得到以下这些索引:

[1, 5]

我该怎么在numpy中实现这个呢?我试过用下面的方式(这是一个快速参考指南上提到的,比较了numpy、matlab和IDL的用法):

(a == 1 and b == 0).nonzero()

但是这样做时出现了一个错误,提示说真值有歧义。

有没有什么好的建议呢?

4 个回答

2

试试下面的代码:

import numpy

a = numpy.array([0, 1, 3, 5, 1, 1, 2])
b = numpy.array([1, 0, 2, 5, 3, 0, 6])

res =  [i for i,v in enumerate(zip(a,b)) if v == (1,0)]

print res
6

这里有一种方法:

In [75]: import numpy as np

In [76]: a = np.array([0, 1, 3, 5, 1, 1, 2])

In [77]: b = np.array([1, 0, 2, 5, 3, 0, 6])

In [78]: np.argwhere((a==1) & (b==0)).flatten()
Out[78]: array([1, 5])
14
In []: from numpy import array, logical_and
In []: a= array([0, 1, 3, 5, 1, 1, 2])
In []: b= array([1, 0, 2, 5, 3, 0, 6])
In []: logical_and(a== 1, b== 0).nonzero()[0]
Out[]: array([1, 5])

显然,这样做也可以:

In []: ((a== 1)& (b== 0)).nonzero()[0]
Out[]: array([1, 5])

撰写回答