如何使用索引数组从 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]]

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