numpy 选取数组某一维度上的切片 numpy.compress

numpy.compress(condition, a, axis=None, out=None)

官方文档

Return selected slices of an array along given axis.
When working along a given axis, a slice along that axis is returned in output for each index where condition evaluates to True. When working on a 1-D array, compress is equivalent to extract.

Parameter
condition [1-D array of bools]
Array that selects which entries to return. If len(condition) is less than the size of a along the given axis, then output is truncated to the length of the condition array.

沿着相应的axis进行切片，并提取condition=1对应位置上的切片重组成新的数组

示例:

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

提取矩阵的第0行和第1行

condition=[1,1,0]
np.compress(condition,a,axis=0)

array([[1, 2],
       [3, 4]])

矩阵有3行，但如果condition的长度不到3行，那么condition自动补0

condition=[1,1]#自动补充成[1,1,0]
np.compress(condition,a,axis=0)

array([[1, 2],
       [3, 4]])

用True,False代替0,1也可

np.compress([True, True, False], a, axis=0)

array([[1, 2],
       [3, 4]])