numpy multiple boolean index arrays

Question

I have an array which I want to use boolean indexing on, with multiple index arrays, each producing a different array. Example:

w = np.array([1,2,3])
b = np.array([[False, True, True], [True, False, False]])

Should return something along the lines of:

[[2,3], [1]]

I assume that since the number of cells containing True can vary between masks, I cannot expect the result to reside in a 2d numpy array, but I'm still hoping for something more elegant than iterating over the masks the appending the result of indexing w by the i-th b mask to it.

Am I missing a better option?

Edit: The next step I want to do afterwards is to sum each of the arrays returned by w[b], returning a list of scalars. If that somehow makes the problem easier, I'd love to know as well.

Are you looking for elegance or performance? Those two aren't the same. — Divakar
– Divakar, Commented Nov 30, 2019 at 17:27
If you are looking to just get the summations, simply use matrix-multiplication - b.dot(w). — Divakar
– Divakar, Commented Nov 30, 2019 at 17:30
I would use dot product (I'm imitating that), but in my real example both b and w are very large and sparse, where b is known to only contain 1s. I'm going for performance here. — Yiftach
– Yiftach, Commented Nov 30, 2019 at 17:32
If those are stored in regular arrays and not sparse matrices, the dot product should still be pretty fast. — Divakar
– Divakar, Commented Nov 30, 2019 at 17:36

Mark · Accepted Answer · 2019-11-30 17:37:29Z

1

Assuming you want a list of numpy arrays you can simply use a comprehension:

w = np.array([1,2,3])
b = np.array([[False, True, True], [True, False, False]])
[w[bool] for bool in b]

# [array([2, 3]), array([1])]

If your goal is just a sum of the masked values you use:

np.sum(w*b) # 6

or

np.sum(w*b, axis=1) # array([5, 1])
# or b @ w

…since False times you number will be 0 and therefor won't effect the sum.

edited Nov 30, 2019 at 17:37

answered Nov 30, 2019 at 17:30

Mark

92.6k8 gold badges116 silver badges156 bronze badges

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

Comments

Bill Chen · Accepted Answer · 2019-11-30 17:33:34Z

0

Try this:

[w[x] for x in b]

Hope this helps.

answered Nov 30, 2019 at 17:33

Bill Chen

1,73914 silver badges24 bronze badges

Collectives™ on Stack Overflow

numpy multiple boolean index arrays

2 Answers 2

Comments

Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Related