python - How to give multiple conditions to numpy.where()

Question

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?

question from:https://stackoverflow.com/questions/65844348/how-to-give-multiple-conditions-to-numpy-where

Answer 1

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]