Python 从3D创建4D上部对角线阵列

假设我有一个

（x，y，z）

大小的矩阵。现在，我想创建一个新的维度矩阵

（x，y，I，I）

，其中

（I，I）

矩阵是上对角线，并由

z

维度上的值构造而成。在

numpy

中是否有一些简单的方法可以做到这一点，而不必使用超过1个for循环（在x上循环）？谢谢

编辑

original = np.array([
    [
        [0, 1, 3],
        [4, 5, 6]
    ],
    [
        [7, 8, 9],
        [3, 2, 1]
    ],
])

new = np.array([
    [
        [
            [0, 1],
            [0, 3]
        ],
        [
            [4, 5],
            [0, 6]
        ]
    ],
    [
        [
            [7, 8],
            [0, 9]
        ],
        [
            [3, 2],
            [0, 1]
        ]
    ]
])

因此，使用上面的方法我们可以看到

original[0, 0, :] = [0 1 3]

new[0, 0, :, :] = [[0 1]
                  [0 3]]

---

下面是一种使用

布尔索引的方法-

n = 2 # This would depend on a.shape[-1]
out = np.zeros(a.shape[:2] + (n,n,),dtype=a.dtype)
out[:,:,np.arange(n)[:,None] <= np.arange(n)] = a

可以建议使用
下标索引
替代最后一步-

r,c = np.triu_indices(n)
out[:,:,r,c] = a

注意：如前所述，
n
取决于
a.shape[-1]
。在这里，我们有
a.shape[-1]
作为
3
，所以
n
是
2
。如果
a.shape[-1]
为
6
，
n
为
3
，依此类推。关系是：
（n*（n+1））//2==a.shape[-1]
您能展示一下
输出[0,0，：，：]
在给定的
输入[0,0，：]
下的样子吗？也许可以用一个例子来说明这一点。@Divakar我添加了一个例子-它有意义吗？老实说，没有什么真正的意义。我知道如何在
x
和
y
维度上循环。不过，这将是太慢了。谢谢你，是的。非常感谢Divakar！这是完美的。是的，我通过简单地从第一个整数的和中倒出来查找
n
。我必须补充的是，这个解决方案是如此优雅，它伤害了我的眼睛
r,c = np.triu_indices(n)
out[:,:,r,c] = a