从 2D 数组创建 3D（部分）对角线数组答案

【问题标题】：Create a 3D (partially) diagonal array from a 2D array从 2D 数组创建 3D（部分）对角线数组
【发布时间】：2023-01-24 21:52:44
【问题描述】：

我想问一下如何从 2D 数组有效地生成一个 numpy 3D 数组，每行填充新数组的对角线部分？例如，输入的二维数组是

array([[1, 2],
       [3, 4],
       [5, 6],
       [7, 8]])

我希望输出是

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

       [[3, 0],
        [0, 4]],

       [[5, 0],
        [0, 6]],

       [[7, 0],
        [0, 8]]])

通常，第一维的大小非常大。提前致谢。

【问题讨论】：

您目前的发电情况如何？为什么您认为它的效率不够高？
我使用了一个 for 循环来诊断每一行，但这效率很低，因为通常输入数组包含数千行
np.einsum('jk,kl->jkl', arr, np.eye(2, dtype=int))或arr[...,None] * np.eye(2, dtype=int)[None]

【解决方案1】：

arr = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
res = np.apply_along_axis(np.diag, 1, arr)

【讨论】：

【解决方案2】：

假设 a 输入并使用 unravel_index 的索引：

x, y = np.unravel_index(np.arange(a.size), a.shape)

out = np.zeros(a.shape+(a.shape[-1],), dtype=a.dtype)

out[x, y, y] = a.flat

输出：

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

       [[3, 0],
        [0, 4]],

       [[5, 0],
        [0, 6]],

       [[7, 0],
        [0, 8]]])

【讨论】：