python - numpy 索引 : fixed length parts of each row with varying starting column

Question

我有一个像 z 这样的二维数组和一个表示“起始列位置”的一维数组，例如 starts。另外我有一个固定的 row_length = 2

z = np.arange(35).reshape(5, -1)
# --> 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, 25, 26, 27],
             [28, 29, 30, 31, 32, 33, 34]])

starts = np.array([1,5,3,3,2])

我想要的是这个缓慢的 for 循环的结果，如果可能的话更快。

result = np.zeros(
    (z.shape[0], row_length), 
    dtype=z.dtype
)
for i in range(z.shape[0]):
    s = starts[i]
    result[i] = z[i, s:s+row_length]

因此，此示例中的结果最终应如下所示:

array([[ 1,  2],
       [12, 13],
       [17, 18],
       [24, 25],
       [30, 31]])

我似乎找不到使用花哨的索引或 np.take 来提供此结果的方法。

Answer 1

一种方法是使用 broadcasted additions 获取这些索引与这些 starts 和 row_length 然后使用 NumPy's advanced-indexing从数据数组中提取所有这些元素，就像这样 -

idx = starts[:,None] + np.arange(row_length)
out = z[np.arange(idx.shape[0])[:,None], idx]

示例运行 -

In [197]: z
Out[197]: 
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, 25, 26, 27],
       [28, 29, 30, 31, 32, 33, 34]])

In [198]: starts = np.array([1,5,3,3,2])

In [199]: row_length = 2

In [200]: idx = starts[:,None] + np.arange(row_length)

In [202]: z[np.arange(idx.shape[0])[:,None], idx]
Out[202]: 
array([[ 1,  2],
       [12, 13],
       [17, 18],
       [24, 25],
       [30, 31]])