如何将二维NumPy数组转换为带0填充的对角矩阵？

问题示例 #

输入二维NumPy数组：

array([[1, 1],
       [2, 2],
       [3, 3]])

需要转换为：

array([[1, 1, 0, 0],
       [0, 2, 2, 0],
       [0, 0, 3, 3]])

np.diag仅适用于一维数组转对角矩阵，循环切片的方法可行但不够高效，以下是更优的向量化实现方案：

高效实现方案 #

利用NumPy的向量化索引操作，避免显式循环，性能更优：

import numpy as np

# 定义输入数组
input_arr = np.array([[1, 1],
                      [2, 2],
                      [3, 3]])

rows, cols = input_arr.shape
# 创建目标形状的零数组
output_arr = np.zeros((rows, rows * cols), dtype=input_arr.dtype)

# 计算每个分块的起始列索引
block_starts = np.arange(rows) * cols
# 通过广播索引填充对应位置
output_arr[np.arange(rows)[:, None], block_starts[:, None] + np.arange(cols)] = input_arr

print(output_arr)

代码解释 #

1. 先获取输入数组的行数rows和列数cols，创建行数不变、列数为rows*cols的零数组作为结果容器。
2. block_starts生成每个分块的起始列位置（比如第0个块从0列开始，第1个从2列开始，第2个从4列开始）。
3. 利用np.arange(rows)[:, None]将行索引转换为列向量，结合block_starts[:, None] + np.arange(cols)生成每个元素对应的目标列索引，通过广播一次性完成所有元素的赋值，完全避免循环。

这种方法充分利用了NumPy的向量化特性，在处理大规模数组时相比循环切片有显著的性能提升。

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