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

Viewed 158

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]]
]
2 Answers

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).

You can do both your tasks in one go (a single loop):

for i in range(a.shape[0]):
    a[i,...] = np.triu(a[i,...], k=1).T

The resul is:

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]]])
Related