NumPy中按行号独立平移数组行并适配目标输出形状的实现探讨

这个需求的核心是为输入数组的每一行设置独立的列偏移量（规则为第i行向右偏移i个位置），最终在零初始化的目标矩阵上填充原数组元素，完全可以依靠NumPy的广播机制实现无硬编码、通用化的优雅写法，比手动拼接索引的方式可读性和扩展性都强很多。

import numpy as np

# 输入数组
A = np.array([[0, 1, 2],
              [1, 2, 3],
              [2, 3, 4]], dtype=np.float32)

n_rows, n_cols = A.shape
# 可根据需求调整目标输出的列数，示例场景为容纳最大偏移的最小宽度
target_cols = n_cols + (n_rows - 1)
out = np.zeros((n_rows, target_cols), dtype=A.dtype)

# 核心索引构建与赋值
i, j = np.ogrid[:n_rows, :n_cols]
out[i, i + j] = A

运行代码得到的out完全匹配示例要求的输出：

array([[0., 1., 2., 0., 0.],
       [0., 1., 2., 3., 0.],
       [0., 0., 2., 3., 4.]], dtype=float32)

• 无硬编码尺寸限制：不管是3x3的小数组，还是10x10、100x100的大规模数组，不需要修改核心逻辑就能直接运行
• 灵活适配任意规则：如果需要调整平移规则，只需要修改列偏移的计算方式即可。比如需要自定义每一行的偏移量，只需要定义和行数长度一致的偏移向量，替换i + j为自定义偏移量 + j，反向偏移、固定偏移等场景都不需要改动其他代码
• 执行效率与手动拼接索引一致：全程依靠NumPy原生的广播机制实现，没有额外的循环或冗余计算，速度和手写索引方式没有差异
• 可自由适配目标输出形状：只需要调整target_cols的数值即可，只要目标宽度不小于「最大行偏移量 + 原列数」就不会出现索引越界的问题

对比目前手写np.r_拼接索引的实现方式，这种写法不需要手动计算切片范围、拼接重复维度，避免了数组尺寸变化时要手动修改硬编码数字的麻烦。

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