使用NumPy构造分块对角矩阵 沿对角线扩展生成同值方块矩阵

核心实现思路 #

利用NumPy广播特性生成块掩码，全向量化操作无Python层面循环，性能远高于手动切片循环。

分输入场景实现 #

场景1：输入为index和counts数组 #

import numpy as np

index = np.array([1, 2, 3])
counts = np.array([2, 3, 2])

# 计算矩阵总边长
total_size = counts.sum()
# 计算每个位置所属的块编号
cum_counts = np.cumsum(counts)
block_idx = np.searchsorted(cum_counts, np.arange(total_size), side='right')
# 广播生成同块掩码，赋值对应数值
mask = block_idx[:, np.newaxis] == block_idx[np.newaxis, :]
result = np.where(mask, index[block_idx], 0)

场景2：输入为展开的arr数组 #

如果已经得到展开的一维数组，实现更简单：

import numpy as np

arr = np.array([1, 1, 2, 2, 2, 3, 3])
mask = arr[:, np.newaxis] == arr[np.newaxis, :]
result = np.where(mask, arr[:, np.newaxis], 0)

可选补充方案（允许使用scipy时） #

如果可以引入scipy依赖，用block_diag接口写法更直观：

import numpy as np
from scipy.linalg import block_diag

index = np.array([1, 2, 3])
counts = np.array([2, 3, 2])
blocks = [val * np.ones((n, n), dtype=int) for val, n in zip(index, counts)]
result = block_diag(*blocks)

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