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

Question

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?

Chrysophylaxs · Accepted Answer · 2022-11-26 13:41:25Z

1

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 &.

edited Nov 26, 2022 at 13:41

answered Nov 26, 2022 at 12:59

Chrysophylaxs

6,5933 gold badges13 silver badges25 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Daniil Yefimov Over a year ago

Thank you for your answer! this works, how can it be done in general(for n columns)? For example if I have a dictionary with n key_value pairs?

Chrysophylaxs Over a year ago

See the edited solution. The evidence being a dictionary feels a bit awkward to me though. If that is something you can change, I would

Collectives™ on Stack Overflow

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

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related