如何将numpy数组每列转为2×2对角块，得到(N,2,2)形状的数组

存在完全向量化的高效实现方案，全程由Numpy底层C语言实现运算，无Python层面循环，即使N取值很大也能保持极高性能：

核心实现代码 #

import numpy as np

# 示例输入数组
arr = np.array([[1, 2, 3],
                [4, 5, 6]])

# 初始化形状为(N,2,2)的全零数组，dtype和输入保持一致
out = np.zeros((arr.shape[1], 2, 2), dtype=arr.dtype)
# 给所有对角块的对角线位置批量赋值
out[:, [0, 1], [0, 1]] = arr.T

print(out)

运行上述代码即可得到预期输出：

[[[1 0]
  [0 4]]

 [[2 0]
  [0 5]]

 [[3 0]
  [0 6]]]

实现逻辑说明 #

• 输入数组形状为(2, N)，转置后得到形状为(N, 2)的数组，每一行对应原数组的一列，也就是每个2×2对角块需要的两个对角元素
• 高级索引[:, [0, 1], [0, 1]]会同时选中所有N个对角块的(0,0)和(1,1)两个对角位置，直接批量赋值即可完成转换

性能对比 #

如果用遍历+np.diag的方式实现（如下所示），当N大于1000时性能会比上述方案低1~2个数量级，仅适合小数据量场景：

# 低性能遍历实现，不推荐
out = np.array([np.diag(col) for col in arr.T])

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