python - 选择行，使特定列包含列表中的值

Question

我想从 numpy 数组中获取在一列中具有特定值的行。下面的例子展示了我的方法。当我给出特定值时，我能够获得行，但是当我将多个值作为 ('4', '8') 时，我没有收到预期的行。

import numpy as np

arr = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]])
arr2 = arr[arr[:,3] == 4]
arr3 = arr[arr[:,3] in ('4', '8')]

arr 是 array([[ 1, 2, 3, 4],[ 5, 6, 7, 8], [ 9, 10, 11, 12]])

arr2 是array([[1, 2, 3, 4]])

arr3 是array([1, 2, 3, 4])。

我应该使用什么命令来获取输出数组([[1, 2, 3, 4], [ 5, 6, 7, 8]]) ?

Answer 1

使用np.in1d为我们正在搜索的元素的任何出现创建掩码，然后只需使用 boolean indexing从输入数组中选择有效行 -

arr[np.in1d(arr[:,3], [4,8])]

sample 运行-

In [149]: arr
Out[149]: 
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])

In [150]: np.in1d(arr[:,3], [4,8]) # Mask of valid ones
Out[150]: array([ True,  True, False], dtype=bool)

In [151]: arr[np.in1d(arr[:,3], [4,8])] # Select rows off arr
Out[151]: 
array([[1, 2, 3, 4],
       [5, 6, 7, 8]])