使用多个值搜索Numpy数组

Question

我有numpy 2d数组有重复值.

我正在搜索这样的数组.

In [104]: import numpy as np

In [105]: array = np.array

In [106]: a = array([[1, 2, 3],
     ...:            [1, 2, 3],
     ...:            [2, 5, 6],
     ...:            [3, 8, 9],
     ...:            [4, 8, 9],
     ...:            [4, 2, 3],
     ...:            [5, 2, 3])

In [107]: num_list = [1, 4, 5]

In [108]: for i in num_list :
     ...:     print(a[np.where(a[:,0] == num_list)])
     ...:
 [[1 2 3]
 [1 2 3]]
[[4 8 9]
 [4 2 3]]
[[5 2 3]]

输入是具有与列0值类似的数字的列表.我想要的最终结果是任何格式的结果行,例如array,list或tuple

array([[1, 2, 3],
       [1, 2, 3],
       [4, 8, 9],
       [4, 2, 3],
       [5, 2, 3]])

我的代码工作正常,但似乎并不像pythonic.有多个值的搜索策略是否更好？

比如a[np.where(a[:,0] == l)]只进行一次查找以获取所有值.

我真正的阵列很大

Answer 1

方法#1:使用np.in1d-

a[np.in1d(a[:,0], num_list)]

方法#2:使用np.searchsorted-

num_arr = np.sort(num_list) # Sort num_list and get as array

# Get indices of occurrences of first column in num_list
idx = np.searchsorted(num_arr, a[:,0])

# Take care of out of bounds cases
idx[idx==len(num_arr)] = 0 

out = a[a[:,0] == num_arr[idx]]