如何给 numpy.where() 提供多个条件

Question

我有一个像这样的 numpy 数组：

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

我正在尝试返回另一个数组，其中包含上面数组中某些项目的所有索引。 我努力了：

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

因此，如果字母是 A 或 C，则索引应存储在数组letter_indexes中

Output 应该是：

0, 2, 3, 5, 6, 8

但事实并非如此，我收到一个错误： TypeError: only integer scalar arrays can be converted to a scalar index

我可以给点建议吗？

Answer 1

您可以使用np.in1d ：

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

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

如果你有多个更复杂的条件，你可以使用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)

Answer 2

numpy.any将返回True / False ，因此不适合在这里使用。

您可以只numpy.where ，但在条件中稍作调整 - 您将按位 OR 运算符引入条件中，如下所示：

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

事实上， numpy也内置了这个功能，如numpy.bitwise_or ：

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

这是一个简短的工作示例：

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]

Answer 3

这是一个可能的解决方案：

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