如何高效向量化接收索引输入并返回2x2 NumPy数组的函数？

假设有函数f(i,j,a)，它接收二维NumPy数组a的行索引i、列索引j，返回以(i,j)为左上角的2x2 NumPy数组，代码示例如下：

import numpy as np

def f(i,j,a):
    b = a[np.arange(2)[:, np.newaxis] + i, np.arange(2) + j]
    return b

（上述代码可通过替换2实现通用化）。我希望将该过程向量化，定义接收i_array、j_array的f(i_array,j_array,a)，本质是遍历i、j。目前用列表推导式可简单实现：

b_array = np.array([f(i,j,a) for i in i_array for j in j_array])

请问是否存在更高效的NumPy函数/索引方法直接实现，还是列表推导式已是最优方案？我知道np.vectorize可实现，但它并不比简单列表推导式更快。

有更高效的纯NumPy向量化实现方案，比列表推导式性能更优，尤其是当i_array和j_array规模较大时。核心是利用广播机制和高级索引一次性生成所有目标子数组，避免Python层面的循环开销。

通用实现（支持任意k×k子数组） #

假设要提取的子数组大小为k×k，步骤如下：

1. 生成子数组内的相对索引网格，调整形状以支持广播：

   k = 2
   # 行偏移：形状(k, k, 1, 1)
   row_offsets = np.arange(k)[:, np.newaxis, np.newaxis, np.newaxis]
   # 列偏移：形状(k, k, 1, 1)
   col_offsets = np.arange(k)[np.newaxis, :, np.newaxis, np.newaxis]
   

2. 将输入的索引数组转换为适合广播的形状：

   # i_array转为(1, 1, m, 1)，m为i_array长度；j_array转为(1, 1, 1, n)，n为j_array长度
   i_grid = i_array[np.newaxis, np.newaxis, :, np.newaxis]
   j_grid = j_array[np.newaxis, np.newaxis, np.newaxis, :]
   

3. 计算所有子数组的绝对索引并提取，最后调整维度顺序匹配列表推导式的输出：

   all_rows = row_offsets + i_grid
   all_cols = col_offsets + j_grid
   # 输出形状为(m, n, k, k)，和列表推导式的结果维度一致
   b_array = a[all_rows, all_cols].transpose(2, 3, 0, 1)
   

简化版（固定2×2子数组） #

针对题目中的2×2场景，可直接封装成函数：

def f_vectorized(i_array, j_array, a):
    row_off = np.arange(2)[:, np.newaxis, np.newaxis, np.newaxis]
    col_off = np.arange(2)[np.newaxis, :, np.newaxis, np.newaxis]
    i_grid = i_array[np.newaxis, np.newaxis, :, np.newaxis]
    j_grid = j_array[np.newaxis, np.newaxis, np.newaxis, :]
    return a[row_off + i_grid, col_off + j_grid].transpose(2, 3, 0, 1)

性能优势 #

这种纯向量化方法完全在NumPy的C底层执行，没有Python循环的额外开销。当i_array和j_array的长度超过几十时，性能会显著优于列表推导式——比如当两者各有100个元素时，向量化方法的速度通常是列表推导式的5-10倍（具体取决于原数组a的规模）。

注意事项 #

• 需确保所有i + 2 <= a.shape[0]和j + 2 <= a.shape[1]，避免索引越界；
• 若需要输出形状为(len(i_array)*len(j_array), 2, 2)（和列表推导式的np.array([...])结果完全一致），可在最后添加.reshape(-1, 2, 2)；
• 该方法同样适用于更大的k值，只需替换代码中的2为目标子数组边长即可。

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