NumPy数组中满足值和索引条件的元素索引

时间:2014-08-22 01:40:09

标签: python numpy

我有一个NumPy数组A。我想知道A中元素的索引等于一个值,哪些索引满足某些条件:

import numpy as np
A = np.array([1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4])

value = 2
ind = np.array([0, 1, 5, 10])  # Index belongs to ind

这是我做的:

B = np.where(A==value)[0]  # Gives the indexes in A for which value = 2
print(B)

[1 5 9]

mask = np.in1d(B, ind)  # Gives the index values that belong to the ind array
print(mask)
array([ True, True, False], dtype=bool)

print B[mask]  # Is the solution
[1 5]

解决方案有效,但我发现它很复杂。此外,in1d执行的操作很慢。有没有更好的方法来实现这一目标?

5 个答案:

答案 0 :(得分:14)

如果您翻转操作指令,可以在一行中执行:

B = ind[A[ind]==value]
print B
[1 5]

打破这一点:

#subselect first
print A[ind]
[1 2 2 3]

#create a mask for the indices
print A[ind]==value
[False  True  True False]

print ind
[ 0  1  5 10]

print ind[A[ind]==value]
[1 5]

答案 1 :(得分:7)

B = np.where(A==value)[0]  #gives the indexes in A for which value = 2
print np.intersect1d(B, ind)
[1 5]

答案 2 :(得分:2)

后两个步骤可以由intersect1D替换。可能也做了一种。除非你能保证你的ind阵列是有序的,否则你不知道如何避免这种情况。

答案 3 :(得分:1)

如何将np.where推迟到最后,如下:

res = (A == value)
mask = np.zeros(A.size)
mask[ind] = 1
print np.where(res * z)[0]

那不应该要求任何排序。

答案 4 :(得分:0)

这有点不同 - 我没有做任何计时测试。

>>> 
>>> A = np.array([1,2,3,4,1,2,3,4,1,2,3,4])
>>> ind = np.array([0,1,5,10])
>>> b = np.ix_(A==2)
>>> np.intersect1d(ind, *b)
array([1, 5])
>>> 

虽然在查看@ Robb的解决方案之后,这可能就是这样做的。