如何使用NumPy按行列索引将二维数组拆分为带NaN填充的3×3子数组

问题说明 #

现有通过np.arange生成的6×6数值二维数组，定义如下：

import numpy as np
big = np.arange(6*6).reshape([6,6])
'''
数组内容：
[[ 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]]
'''

需要将该数组拆分为与原数组每个元素一一对应的3×3小型二维数组，规则如下：

• 每个3×3小数组的中心位置值，对应原数组同行列坐标的元素值
• 超出原数组边界的位置（即示例中标注为n的位置）填充NaN
• 小数组按原数组的行列顺序排布：原数组第一行元素对应第一组6个3×3数组，原数组最后一行元素对应最后一组6个3×3数组
  排布规律和示例一致：
• 原数组第一行（元素0~5）对应的3×3数组，最顶行全为NaN
• 原数组最后一行（元素30~35）对应的3×3数组，最底行全为NaN
• 原数组第一列元素对应的3×3数组，最左列全为NaN
• 原数组最后一列元素对应的3×3数组，最右列全为NaN

实现代码 #

直接用NumPy内置的边界填充+滑动窗口功能即可实现，无需手动写循环：

# 1. 将原数组转为浮点类型（整数类型不支持NaN值），四周各填充1层NaN
padded_arr = np.pad(
    big.astype(np.float64),
    pad_width=1,
    mode="constant",
    constant_values=np.nan
)
# 2. 生成滑动窗口视图，直接得到所有3×3子数组
# 输出数组形状为(6, 6, 3, 3)，维度含义为：(原数组行索引, 原数组列索引, 子数组行索引, 子数组列索引)
from numpy.lib.stride_tricks import sliding_window_view
all_subarrays = sliding_window_view(padded_arr, (3, 3))

结果验证 #

取几个典型位置的子数组核对，完全匹配示例要求：

• 原数组[0,0]位置（值为0）对应的子数组：

print(all_subarrays[0, 0])
'''
输出：
[[nan nan nan]
 [nan  0.  1.]
 [nan  6.  7.]]
'''

• 原数组[1,1]位置（值为7）对应的子数组：

print(all_subarrays[1, 1])
'''
输出：
[[ 0.  1.  2.]
 [ 6.  7.  8.]
 [12. 13. 14.]]
'''

• 原数组[5,5]位置（值为35）对应的子数组：

print(all_subarrays[5, 5])
'''
输出：
[[28. 29. nan]
 [34. 35. nan]
 [nan nan nan]]
'''

该实现为纯向量化操作，即使原数组尺寸很大也能保持极高的运行效率。

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