Numpy从每个numpy数组行中选择不同数量的第一个元素

Question

我想从矩阵中选择不同数量的第一个元素.数字在数组中指定.结果是一维数组.例如:

a = np.arange(25).reshape([5, 5])
numbers = np.array([3, 2, 0, 1, 2])

我想要这个结果:

[0, 1, 2, 5, 6, 15, 20, 21]

没有for循环.

Answer 1

让我们用一些NumPy broadcasting魔法!

a[numbers[:,None] > np.arange(a.shape[1])]

样品运行 -

In [161]: a
Out[161]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

In [162]: numbers
Out[162]: array([3, 2, 0, 1, 2])

In [163]: numbers[:,None] > np.arange(a.shape[1]) # Mask to select elems
Out[163]: 
array([[ True,  True,  True, False, False],
       [ True,  True, False, False, False],
       [False, False, False, False, False],
       [ True, False, False, False, False],
       [ True,  True, False, False, False]], dtype=bool)

In [164]: a[numbers[:,None] > np.arange(a.shape[1])] # Select w/ boolean indexing
Out[164]: array([ 0,  1,  2,  5,  6, 15, 20, 21])