Selecting rows in 2-D numpy array based on subset of column values-CodePudding

Suppose I have the following numpy array:

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]])

I want to select only the rows, where column with index 1 have value 1 and column with index 2 have value 0.

i tried the following:

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

But after that i am stuck. how can I do it in numpy 2-D array?

CodePudding user response：

Try:

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

Given a dictionary of column, value pairs, you could use:

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

which generalizes the above solution to a sequence of &.