基于N维数组构造(N+1)维对角矩阵的无循环高效实现

嘿，我完全懂你这种不想用显式循环处理大规模数组的痛点！numpy其实有两个非常高效的方案，完全依赖内置的矢量操作和广播/高级索引，性能拉满，根本不用写循环。

方案一：广播 + 单位矩阵相乘 #

这种方法最简洁，利用numpy的广播机制自动完成元素映射：

假设你的原数组是arr，最后一维的长度为D = arr.shape[-1]，我们只需要：

import numpy as np

# 示例原数组：shape (5,3)
arr = np.random.rand(5, 3)
# 扩展为(5,3,3)的目标数组
expanded_arr = arr[..., np.newaxis] * np.eye(arr.shape[-1])

原理解释： #

• np.eye(arr.shape[-1])生成一个D×D的单位对角矩阵，所有对角位置为1，其余为0。
• arr[..., np.newaxis]给原数组最后添加一个维度，把shape从(*original_shape, D)变成(*original_shape, D, 1)。
• 利用numpy的广播特性，这两个数组会自动匹配维度：单位矩阵被广播为(*original_shape, D, D)，原数组扩展后的版本被广播为(*original_shape, D, D)，相乘后原数组的每个元素就刚好落在目标数组的对角位置上，非对角位置保持0。

方案二：高级索引直接赋值 #

如果你需要自定义非对角位置的初始值（比如不想用0填充），可以先创建目标数组，再用高级索引把原数组元素放到对角位置：

# 创建目标数组，shape为原数组shape + (D,)，这里用1初始化示例
expanded_arr = np.ones(arr.shape + (arr.shape[-1],), dtype=arr.dtype)
# 获取对角位置的索引
idx = np.arange(arr.shape[-1])
# 把原数组元素赋值到对角位置
expanded_arr[..., idx, idx] = arr

原理解释： #

• expanded_arr[..., idx, idx]是numpy的高级索引语法：...匹配前面所有维度，idx, idx表示在最后两个维度上取(0,0), (1,1), ..., (D-1,D-1)的位置，刚好就是对角位置。
• 这种方式直接通过矢量索引赋值，完全没有循环，性能和numpy内置操作一致。

适用场景 #

这两种方法都适用于任意N维数组：不管原数组是2维、3维还是更高维，只要替换arr即可自动适配，比如原数组shape为(a,b,c,d)，扩展后会变成(a,b,c,d,d)，完美满足你的需求。

性能对比 #

这两种方法都是numpy底层优化的操作，比显式循环、列表推导甚至map快几个数量级，处理大规模数组（比如百万级元素）的时候优势特别明显。

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