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

2024-09-28 22:22:35 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个NumPy数组,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做的排序也很慢。有没有更好的方法来实现这个目标?


Tags: thetonumpytrue元素valuenpmask
3条回答
B = np.where(A==value)[0]  #gives the indexes in A for which value = 2
print np.intersect1d(B, ind)
[1 5]

如果将操作顺序翻转过来,可以在一行中完成:

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]

后两个步骤可以用intersect1D代替。可能也有点问题。不知道如何避免,除非你能保证你的ind数组是有序的。

相关问题 更多 >