如何无需显式循环，基于多偏移量从二维数组索引多个子块

假设有一个10×10的二维数组A：

# 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, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
       [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
       [50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
       [70, 71, 72, 73, 74, 75, 76, 77, 78, 79],
       [80, 81, 82, 83, 84, 85, 86, 87, 88, 89],
       [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])

同时有一个3×2的偏移量数组offsets，以及block_size=3：

# offsets
array([[0, 0],
       [2, 3],
       [4, 5]])

我需要从A中按offsets指定的起始位置，提取尺寸为block_size×block_size的子块，目前用显式for循环可以实现需求：

block_size = 3
results = []
for x, y in offsets:
    block = A[x:x+block_size, y:y+block_size]
    results.append(block)
results = np.array(results)

# results.shape = (3, 3, 3)
array([[[ 0,  1,  2],
        [10, 11, 12],
        [20, 21, 22]],

       [[23, 24, 25],
        [33, 34, 35],
        [43, 44, 45]],

       [[45, 46, 47],
        [55, 56, 57],
        [65, 66, 67]]])

请问有没有不需要显式for循环的实现方式？比如类似下面这种（但写法无效）的方法：

# 该写法无效
results = A[offsets[:,0]:offsets[:,0]+block_size, offsets[:,1]:offsets[:,1]+block_size]

可以用NumPy的向量化操作实现，这里提供几种高效的无循环方案：

方法1：广播生成索引矩阵 #

通过广播生成所有子块的行、列索引，再直接索引数组A：

import numpy as np

block_size = 3
# 生成块内的偏移量：(block_size, block_size)
row_offs = np.arange(block_size)[:, None]
col_offs = np.arange(block_size)
# 扩展offsets维度，与块内偏移量广播：(3, block_size, block_size)
rows = offsets[:, 0, None, None] + row_offs
cols = offsets[:, 1, None, None] + col_offs
# 索引得到结果
results = A[rows, cols]

得到的results形状为(3, 3, 3)，和循环结果完全一致。

方法2：使用np.lib.stride_tricks.as_strided（内存高效） #

如果数组A是连续内存的，可以用as_strided直接生成视图，避免数据复制：

from numpy.lib.stride_tricks import as_strided

block_size = 3
# 获取原数组的内存步长
strides = A.strides
# 生成所有可能的block视图，再按offsets提取目标块
all_blocks = as_strided(A, 
                        shape=(A.shape[0]-block_size+1, A.shape[1]-block_size+1, block_size, block_size),
                        strides=(strides[0], strides[1], strides[0], strides[1]))
results = all_blocks[offsets[:,0], offsets[:,1]]

注意：这种方法生成的是原数组的视图，修改results会直接影响原数组A，若需要独立数组可添加.copy()。

方法3：使用np.take_along_axis #

结合广播索引和take_along_axis分步提取：

block_size = 3
# 生成每个子块的行索引范围：(3, block_size)
row_indices = offsets[:,0, None] + np.arange(block_size)
# 生成每个子块的列索引范围：(3, block_size)
col_indices = offsets[:,1, None] + np.arange(block_size)
# 先按行提取对应行，再按列提取对应列
temp = np.take_along_axis(A, row_indices[:, :, None], axis=0)
results = np.take_along_axis(temp, col_indices[:, None, :], axis=1)

这种方法逻辑直观，同样能得到符合要求的结果。

内容的提问来源于stack exchange，提问作者Qimin Chen