我有一个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
会进行一个比较慢的排序。有没有更好的方法来实现这个功能?