'How to get upper triangles of an array of 2d-arrays

Let's say I have an array that contains 2 times 3x3-array:

a = np.arange(2 * 3 * 3).reshape(2, 3, 3)
print(a)

Output:

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

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]]])

Now I would like to have the upper triangle of each of the 2 arrays. I know I can achieve it through the following two:

np.array([aa[np.triu_indices(3)] for aa in a])
# or 
a.T[np.tril_indices(3)].T

Output:

array([[ 0,  1,  2,  4,  5,  8],
       [ 9, 10, 11, 13, 14, 17]])

However, I know that list comprehension is slow so I'd rather not use it. And the transpose + tril makes it difficult to understand what it does at first sight. I had hoped that one of the following options would work, but none of them did:

a[:, np.triu_indices(3)] # totally different output
a[np.arange(len(a)), np.triu_indices(3)] # error
a[np.indices(a.shape)[0], np.triu_indices(3)] # error

Is there an elegant and fast way to do it?



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source