在给定布尔索引的情况下重新调整 numpy 数组中的列

Question

我有给定的数据集：

data = np.array([
    [1, 2, 1, 3, 1, 2, 1],
    [3, 4, 1, 5, 2, 7, 2],
    [2, 1, 2, 1, 1, 4, 5],
    [6, 1, 2 ,3, 1, 3, 1]])

cols_idx = np.array([0, 0, 1, 0, 1, 0, 0])

我想从cols_idx == 1 data返回列。 为此，我使用了：

data[:, np.nonzero(cols_idx)]

但它返回一个 3D 而不是 2D 数组：

data[:, np.nonzero(cols_idx)]
array([[[1, 1]],    
       [[1, 2]],
       [[2, 1]],    
       [[2, 1]]])

data[:, np.nonzero(cols_idx)].shape
(4, 1, 2)

我希望输出是：

data[:, np.nonzero(cols_idx)]
array([[1, 1],    
       [1, 2],
       [2, 1],    
       [2, 1]])

data[:, np.nonzero(cols_idx)].shape
(4, 2)

我怎样才能做到这一点？

Answer 1

print(np.nonzero(cols_idx))给出(array([2, 4]),) （一个元组而不仅仅是一个数组）

所以你应该使用np.nonzero(cols_idx)[0] # gives [2 4]来得到你想要的：

完整代码：

import numpy as np 
data = np.array([
    [1, 2, 1, 3, 1, 2, 1],
    [3, 4, 1, 5, 2, 7, 2],
    [2, 1, 2, 1, 1, 4, 5],
    [6, 1, 2 ,3, 1, 3, 1]])

cols_idx = np.array([0, 0, 1, 0, 1, 0, 0])
new_data = data[:, np.nonzero(cols_idx)[0]]
print(new_data)
'''[[1 1]                                                                                                                        
 [1 2]                                                                                                                        
 [2 1]                                                                                                                        
 [2 1]]'''
print(new_data.shape) # (4,2)

Answer 2

从numpy文档：

虽然可以使用 a[nonzero(a)] 获得非零值，但建议使用 x[x.astype(bool)] 或 x[x != 0] 代替，这将正确处理 0-d 数组。

所以最好使用：

data[:, cols_idx.astype(bool)]

或者

data[:, cols_idx != 0]