使用1d布尔索引从2d数组中选择 [英] Use 1d boolean index to select out of 2d array

问题描述

有时候我会有一个ND数组，我需要从中选择数据，但是数据标准只有M < N维.举个例子

Sometimes I'll have an ND array out of which I need to select data, but the data criterion has only M < N dimensions. Take for example

## 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:

but that leads to an IndexError:

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

自然地，如果我在第二维上重复布尔索引，则一切正常，以下是预期的输出:

Naturally, if I repeat the boolean index on the second dimension, everything works -- the following is the expected output:

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

但是，这效率很低.用numpy实现此目的而不必重复矩阵的自然方法是什么?

However, this is quite inefficient. What's the natural way of achieving this with numpy without having to repeat the matrix?

推荐答案

如果我了解您的问题，则可以使用省略号(...)来覆盖未过滤的尺寸:

If I understand your problem, you can use ellipsis (...) to cover unfiltered dimensions:

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

这篇关于使用1d布尔索引从2d数组中选择的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！