基于布尔向量选择 numpy 中的列

Selecting columns in numpy based on a Boolean vector

我有两个 NumPy 数组 ab,尺寸为 m,乘以 n。我有一个长度为 n 的布尔向量 b,我想生成一个新数组 c,它从 a、[=11] 中选择 n 列=],所以如果 b[i] 为真,我从 b 中获取列,否则从 a 中获取列。

如何以最有效的方式执行此操作? 我看过 selectwherechoose

首先,让我们设置一些示例代码:

import numpy as np

m, n = 5, 3
a = np.zeros((m, n))
b = np.ones((m, n))

boolvec = np.random.randint(0, 2, m).astype(bool)

只是为了展示这些数据的样子:

In [2]: a
Out[2]: 
array([[ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.]])

In [3]: b
Out[3]: 
array([[ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.]])

In [4]: boolvec
Out[4]: array([ True,  True, False, False, False], dtype=bool)

在这种情况下,使用 np.where 最有效。但是,我们需要 boolvec 的形状可以广播到与 ab 相同的形状。因此,我们可以通过使用 np.newaxisNone (它们相同)进行切片使其成为列向量:

In [5]: boolvec[:,None]
Out[5]: 
array([[ True],
       [ True],
       [False],
       [False],
       [False]], dtype=bool)

然后我们可以使用np.where得到最终结果:

In [6]: c = np.where(boolvec[:, None], a, b)

In [7]: c
Out[7]: 
array([[ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.]])

您可以为此使用 np.choose

例如 ab 数组:

>>> a = np.arange(12).reshape(3,4)
>>> b = np.arange(12).reshape(3,4) + 100
>>> a_and_b = np.array([a, b])

要使用np.choose,我们需要一个包含两个数组的 3D 数组; a_and_b 看起来像这样:

array([[[  0,   1,   2,   3],
        [  4,   5,   6,   7],
        [  8,   9,  10,  11]],

       [[100, 101, 102, 103],
        [104, 105, 106, 107],
        [108, 109, 110, 111]]])

现在让布尔数组为bl = np.array([0, 1, 1, 0])。那么:

>>> np.choose(bl, a_and_b)
array([[  0, 101, 102,   3],
       [  4, 105, 106,   7],
       [  8, 109, 110,  11]])

(5000,3000) 个数组的时间是:

In [107]: timeit np.where(boolvec[:,None],b,a)
1 loops, best of 3: 993 ms per loop

In [108]: timeit np.choose(boolvec[:,None],[a,b])
1 loops, best of 3: 929 ms per loop

In [109]: timeit c=a[:];c[boolvec,:]=b[boolvec,:]
1 loops, best of 3: 786 ms per loop

wherechoose本质上是一样的;布尔索引稍微快一些。 select用的是choose,所以没计时


我的列采样时间与此类似,只是这次索引速度较慢:

In [119]: timeit np.where(cols,b,a)
1 loops, best of 3: 878 ms per loop

In [120]: timeit np.choose(cols,[a,b])
1 loops, best of 3: 915 ms per loop

In [121]: timeit c=a[:];c[:,cols]=b[:,cols]
1 loops, best of 3: 1.25 s per loop

更正,对于索引,我应该使用 a.copy()

In [32]: timeit c=a.copy();c[boolvec,:]=b[boolvec,:]
1 loops, best of 3: 783 ms per loop
In [33]: timeit c=a.copy();c[:,cols]=b[:,cols]
1 loops, best of 3: 1.44 s per loop

Python2.7 和 3、numpy 1.8.2 和 1.9.0 dev

的时间相同