无需循环实现指定格式的矩阵填充优化方法

问题背景 #

给定numpy矩阵：

import numpy as np
A = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])

需要将其转换为每行首尾分别填充该行首、尾元素的结果：

result = np.array([[1., 1., 2., 3., 3.],
                   [4., 4., 5., 6., 6.],
                   [7., 7., 8., 9., 9.]])

原实现通过循环逐行调用np.pad完成，现提供三种无循环的矢量化优化方案：

方案一：横向拼接首尾列与原矩阵 #

利用numpy的切片和hstack直接拼接，逻辑简单高效：

# 提取每行首元素转为列向量
first_col = A[:, [0]]
# 提取每行尾元素转为列向量
last_col = A[:, [-1]]
# 横向拼接得到结果
result = np.hstack([first_col, A, last_col])
# 若需要转为float类型，可添加：
# result = result.astype(float)

这种方式完全基于矢量化操作，没有循环开销，是最直观的实现方式。

方案二：批量调用np.pad #

通过np.pad的constant_values参数传入每行的首尾元素数组，实现批量填充：

result = np.pad(
    A,
    pad_width=((0, 0), (1, 1)),  # 行维度不填充，列维度左右各填1个元素
    mode='constant',
    constant_values=(A[:, 0].reshape(-1, 1), A[:, -1].reshape(-1, 1))
)

这里将首尾元素转为二维数组，保证constant_values的形状能与原矩阵广播匹配，从而自动完成每行的填充。

方案三：索引直接赋值构造结果 #

先创建目标形状的空数组，再通过索引将对应值填入：

# 初始化结果数组，类型与原矩阵一致
result = np.zeros((A.shape[0], A.shape[1] + 2), dtype=A.dtype)
# 中间列填入原矩阵数据
result[:, 1:-1] = A
# 首列填入每行的首元素
result[:, 0] = A[:, 0]
# 尾列填入每行的尾元素
result[:, -1] = A[:, -1]

此方法逻辑清晰，适合需要精细控制数组构造的场景。

以上三种方案均为矢量化操作，相比循环实现，在矩阵规模较大时能显著提升运行效率。

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