CODEWITHSUNDEEP
pythonarraysnumpyindexing
Tom de Geus
Tom de Geus

Reputation: 5975

'where' and slice at the same time

I have a list of logicals test and a list of indices i of the same length. Now I want to select a[i] where test == True. Naturally I use np.where, e.g.

import numpy as np

a = np.arange(4)
test = np.array([True, False, False, True])
i = np.array([0, 0, 0, 0])

b = np.where(test, a[i], -1)
print(b)

But now I encounter the case where the indices in i are not valid indices of a where test == False. E.g.

test = np.array([True, False, False, True])
i = np.array([0, 6, 6, 0])

the above code of course fails because the slice does not work. But since I was not interested in the slice where test == False, how can I avoid its evaluation?

Of course a work-around is

b = np.where(test, a[np.where(test, i, 0)], -1)

but this I want to avoid.

Upvotes: 1

Views: 151

Answers (1)

Ehsan
Ehsan

Reputation: 12407

How about this modification on your solution:

b = np.where(test, a[i*test], -1)

It works even if i is out of bound when test is False. Trick is multiplying index to True will return same index whereas multiplying index to False turn it to 0.

Another easy solution is this:

b = -1*np.ones_like(a)
b[test] = a[i[test]]

Upvotes: 2

Related Questions