各行

Question

のための別の列を選択して、私は次の配列があるとします。

>>> a = np.arange(25).reshape((5, 5)) 
>>> a 
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]]) 

今、私は、次のインデックス配列に基づいて、行ごとに異なる列を選択します：

>>> i = np.array([0, 1, 2, 1, 0]) 

このインデックス配列各行の開始列を示し、選択範囲は類似した範囲でなければなりません。私は複数の列について

>>> a[np.arange(a.shape[0]), i] 

が、どのように経由して、行ごとに単一の列を選択することができることを知って

>>> ??? 
array([[ 0, 1, 2], 
     [ 6, 7, 8], 
     [12, 13, 14], 
     [16, 17, 18], 
     [20, 21, 22]]) 

：3.したがって、私は、次のような結果を得ることがしたいですか？

Answer 1

advanced indexingを適切にブロードキャストされた2dアレイをインデックスとして使用してください。

a[np.arange(a.shape[0])[:,None], i[:,None] + np.arange(3)] 
#array([[ 0, 1, 2], 
#  [ 6, 7, 8], 
#  [12, 13, 14], 
#  [16, 17, 18], 
#  [20, 21, 22]]) 

---

idx_row = np.arange(a.shape[0])[:,None] 
idx_col = i[:,None] + np.arange(3) 

idx_row 
#array([[0], 
#  [1], 
#  [2], 
#  [3], 
#  [4]]) 

idx_col 
#array([[0, 1, 2], 
#  [1, 2, 3], 
#  [2, 3, 4], 
#  [1, 2, 3], 
#  [0, 1, 2]]) 

a[idx_row, idx_col] 
#array([[ 0, 1, 2], 
#  [ 6, 7, 8], 
#  [12, 13, 14], 
#  [16, 17, 18], 
#  [20, 21, 22]])