Numpyは、各numpy配列の行とは異なる最初の要素の数を選択します

Question

行列から最初の要素の数を変えたいとします。 数値は配列で指定されます。 結果は1次元配列です。たとえば ：forループのない

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

：

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

は、そして私は、この結果が欲しいです。

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