python - 使用一维 bool 索引从二维数组中选择

Question

有时我会有一个ND我需要从中选择数据的数组，但数据标准只有 M < N方面。举个例子

## generate some matrix
test = np.arange(9).reshape((3, 3))
## some condition based on first-dimension only
selectMe = np.array([ True,  True, False], dtype=bool)

现在我想做

test[selectMe[:, None]]

但这会导致 IndexError:

IndexError: boolean index did not match indexed array along dimension 1; dimension is 3 but corresponding boolean dimension is 1

当然，如果我在第二个维度上重复 bool 索引，一切都会正常 - 以下是预期的输出:

test[np.repeat(selectMe[:, None], 3, axis=1)]
Out[41]: array([0, 1, 2, 3, 4, 5])

但是，这是相当低效的。使用 numpy 实现这一目标的自然方法是什么？无需重复矩阵？

Answer 1

如果我理解您的问题，您可以使用省略号 (...) 来覆盖未过滤的维度:

import numpy as np

test = np.arange(10000).reshape((100, 100))

# condition
selectMe = np.random.randint(0, 2, 100).astype(bool)

assert (test[selectMe, ...].ravel() == test[np.repeat(selectMe[:, None], 100, axis=1)]).all()

%timeit test[selectMe, ...].ravel()                       # 11.6 µs
%timeit test[np.repeat(selectMe[:, None], 100, axis=1)]   # 103 µs