搜索具有多个值的Numpy数组

发布于 2021-01-29 15:03:58

我有具有重复值的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([[1, 2, 3],
       [1, 2, 3],
       [4, 8, 9],
       [4, 2, 3],
       [5, 2, 3]])

我的代码工作正常,但似乎不是pythonic。有没有更好的多值搜索策略?

就像a[np.where(a[:,0] == l)]只进行一次查找即可获取所有值的地方。

我的真实数组很大

关注者
0
被浏览
48
1 个回答
  • 面试哥
    面试哥 2021-01-29
    为面试而生,有面试问题,就找面试哥。

    方法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]]
    


知识点
面圈网VIP题库

面圈网VIP题库全新上线,海量真题题库资源。 90大类考试,超10万份考试真题开放下载啦

去下载看看