根据numpy矩阵选择元素，具体取决于一组值

Question

我在python中的numpy中有一个矩阵，当第一列中的元素为1或2时，我需要选择第三列中的元素，第二列中的元素为2或4。

import numpy as np 
a = np.matrix('1 2 5; 3 4 3; 2 3 2; 2 4 5; 1 3 2; 3 2 2')

我尝试了此代码，但它不起作用。

l1 = np.array([1, 2])
l2 = np.array([2, 4])
result = a[(a[:,0] in l1) & (a[:,1] in l2)]

Answer 1

您可以使用np.in1d：

l1 = np.array([1, 2])
l2 = np.array([2, 4])
result = a[np.in1d(a[:,0], l1) & np.in1d(a[:,1], l2)]
result
#matrix([[1, 2, 5],
#        [2, 4, 5]])