How to give multiple conditions to numpy.where()

user_12903247 0
Viewed 323

I have a numpy array like this:

letters = np.array([A, B, C, A, B, C, A, B, C])

I'm trying to return another array containing all the indexes of certain items in the array above. I have tried:

letter_indexes = np.where(np.any(letters == 'A', letters == 'C'))

So if letter is either A or C, index should be stored in array letter_indexes

Output should be:

0, 2, 3, 5, 6, 8

But it's not, I'm getting an error: TypeError: only integer scalar arrays can be converted to a scalar index

Could I have some advise?

numpy arrays python
3 Answers

You can use np.in1d:

np.where(np.in1d(letters, ['A', 'C']))[0]

Out[]: array([0, 2, 3, 5, 6, 8], dtype=int64)

If you've got multiple more complex conditions, you can use functools.reduce

from functools import reduce
conditions = [letters == 'A', letters == 'C']
np.where(reduce(np.logical_or, conditions))[0]

Out[]: array([0, 2, 3, 5, 6, 8], dtype=int64)
user_44277770

numpy.any will return True/False, so that won't be suitable to use here.

You can stick to just numpy.where, but with a small tweak in the condition -- you introduce the bitwise OR operator into the condition like this:

np.where((letters == 'A') | (letters == 'C'))

In fact, numpy has this built in as well, as numpy.bitwise_or:

np.where(np.bitwise_or(letters == 'A', letters == 'C'))

Here's a short working example:

import numpy as np

A = 'A'
B = 'B'
C = 'C'

letters = np.array([A, B, C, A, B, C, A, B, C])

# Either this:
letter_indexes = np.where((letters == 'A') | (letters == 'C'))

# Or this:
letter_indexes = np.where(np.bitwise_or(letters == 'A', letters == 'C'))

print(letter_indexes[0])

# [0 2 3 5 6 8]
user_147225620

Here is a possible solution:

np.where(np.logical_or(letters == 'A', letters == 'C'))[0]
user_52961060
Related