Search code examples
pythonpython-3.xnumpynumpy-ndarraynumpy-ufunc

How to make element of 3D array into upper triangular and then tranpose it

For example, I got the 3D array below

[[[1,2,3],
[4,5,6]
[7,8,9]],

[[1,3,5],
[2,4,6],
[5,7,9]]

[[1,4,6],
[2,4,7],
[5,8,9]]
]

The first question is that how I can make each element along the first axis become the triangular matrix, i.e

[[[0,2,3],
[0,0,6]
[0,0,0]],

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

[[0,4,6],
[0,0,7],
[0,0,0]]
]

Based on this, how can I then transpose each of them, like

[[[0,0,0],
[2,0,0]
[3,6,0]],

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

[[0,0,0],
[4,0,0],
[6,7,0]]
]
Solution

  • Use np.triu and then swap axes along last two axes to effectively do transpose -

    In [10]: np.triu(a,1).swapaxes(1,2)
Out[10]: 
array([[[0, 0, 0],
        [2, 0, 0],
        [3, 6, 0]],

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

       [[0, 0, 0],
        [4, 0, 0],
        [6, 7, 0]]])

    Swapping can also be achieved with ndarray.transpose(0,2,1).