根據布爾向量選擇numpy中的列

Question

我有兩個NumPy數組a ， b ，維度為m乘n 。 我有一個長度為n的布爾向量b ，我想生成一個新的數組c ，它從a ， b選擇n列，這樣如果b[i]為真，我從b獲取列，否則從a 。

我如何以最有效的方式做到這一點？ 我看了select ， where choose 。

Answer 1

首先，讓我們設置一些示例代碼：

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的形狀可以廣播到與a和b相同的形狀。 因此，我們可以通過使用np.newaxis或None切片使它成為列向量（它們是相同的）：

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.]])

Answer 2

你可以使用np.choose 。

例如a和b數組：

>>> 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]])

Answer 3

（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

where和choose基本相同; 布爾索引略快一些。 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的相同時序