如何使用索引数组矢量化从 numpy 数组中获取子数组

Question

我想使用某种类型的索引数组（索引数组的样式/格式开放供建议）从基本数组中获取一个 numpy 子数组数组。我可以使用 for 循环轻松做到这一点，但想知道是否有一种巧妙的方法来使用 numpy 广播？

约束：保证子数组大小相同。

up_idx = np.array([[0, 0],
                   [0, 2],
                   [1, 1]])
lw_idx = np.array([[2, 2],
                   [2, 4],
                   [3, 3]])
base = np.array([[1, 2, 3, 4],
                 [5, 6, 7, 8],
                 [9, 10, 11, 12]])

samples = []

for index in range(up_idx.shape[0]):
    up_row = up_idx[index, 0]
    up_col = up_idx[index, 1]
    lw_row = lw_idx[index, 0]
    lw_col = lw_idx[index, 1]

    samples.append(base[up_row:lw_row, up_col:lw_col])

samples = np.array(samples)

print(samples)
> [[[ 1  2]
    [ 5  6]]

   [[ 3  4]
    [ 7  8]]

  [[ 6  7]
   [10 11]]]

我试过了：

vector_s = base[up_idx[:, 0]:lw_idx[:, 1], up_idx[:, 1]:lw_idx[:, 1]]

但这似乎是荒谬的。

Answer 1

我认为通常没有一种快速的方法可以通过 numpy 广播操作来做到这一点 - 一方面，您设置问题的方式无法保证生成的子数组将具有相同的形状，从而能够适应单个输出数组。

解决这个问题的最简洁有效的方法可能是通过列表推导；例如

result = np.array([base[i1:i2, j1:j2] for (i1, j1), (i2, j2) in zip(up_idx, lw_idx)])

除非您的基本数组非常大，否则这应该不是太大的瓶颈。

如果您有不同的问题约束（即每种情况下的切片大小相同），则可能会根据花哨的索引提出更快的矢量化解决方案。例如，如果每个切片的大小为 2（如上面的示例），那么您可以使用这样的精美索引来获得相同的结果：

i, j = up_idx.T[:, :, None] + np.arange(2)
result = base[i[:, :, None], j[:, None]]

理解这种奇特的索引的关键是要意识到结果遵循索引数组的广播形状。