如何高效将二维numpy数组的各对角线转换为补0填充的行

高效实现方案 #

你可以通过numpy矢量化索引实现，完全避免Python层面的循环，性能最优，实现代码如下：

import numpy as np

# 输入数组
A = np.array([[1, 7, 5, 0, 5],
              [9, 1, 4, 6, 0],
              [9, 6, 1, 0, 0],
              [2, 5, 0, 0, 0],
              [1, 0, 0, 0, 0]])

n = A.shape[0]
# 生成列索引矩阵
col_idx = np.arange(n)[np.newaxis, :].repeat(n, axis=0)
# 生成对应A的行索引矩阵
row_idx = np.arange(n)[:, np.newaxis] - col_idx
# 构造有效位置掩码（j <= i的位置需要赋值）
mask = col_idx <= np.arange(n)[:, np.newaxis]

# 初始化B为全0数组
B = np.zeros_like(A)
# 直接批量赋值
B[mask] = A[row_idx[mask], col_idx[mask]]

运行后得到的B就是你需要的结果：

array([[1, 0, 0, 0, 0],
       [9, 7, 0, 0, 0],
       [9, 1, 5, 0, 0],
       [2, 6, 4, 0, 0],
       [1, 5, 1, 6, 5]])

更易理解的实现（适合小数据量） #

如果你的数组规模不大，也可以用提取对角线的方式实现，逻辑更直观：

import numpy as np

A = np.array([[1, 7, 5, 0, 5],
              [9, 1, 4, 6, 0],
              [9, 6, 1, 0, 0],
              [2, 5, 0, 0, 0],
              [1, 0, 0, 0, 0]])

n = A.shape[0]
B = np.zeros_like(A)
# 上下翻转A
A_flip = np.flipud(A)

for k in range(-(n-1), 1):
    diag_vals = np.diag(A_flip, k=k)
    row_idx = n + k - 1
    B[row_idx, :len(diag_vals)] = diag_vals

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