根据列值的子集选择二维 numpy 数组中的行

Question

假设我有以下 numpy 数组：

a = np.array([[1, 1, 0, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 1],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 1, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 0],
       [1, 0, 1, 1, 0]])

我只想选择行，其中索引为 1 的列的值为 1，索引为 2 的列的值为 0。

我尝试了以下内容：

evidence = {1:1,2:0}
mask = a[:,list(evidence.keys())] == list(evidence.values())

但在那之后我被困住了。 我怎样才能在 numpy 二维数组中做到这一点？

Answer 1

尝试：

out = a[(a[:, 1] == 1) & (a[:, 2] == 0)]

给定一个包含列、值对的字典，您可以使用：

evidence = {1:1,2:0}
out = a[ np.logical_and.reduce([a[:, c] == v for c, v in evidence.items()]) ]

将上述解决方案概括为 & 序列。