为每个列选择不同的列

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

此索引数组表示每行的开始列，选择范围应相似，例如3。因此我想得到以下结果：

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

我知道我可以通过

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

但是多列呢？你知道吗

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