How to get triangle upper matrix without the diagonal using numpy

Question

Lets say I have the following matrix:

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

How can I extract the upper triangle matrix without the diagonal efficiently? The output would be the following array:

B = np.array([2,3,6])

Answer 1

Short answer

A[np.triu_indices_from(A, k=1)]

Long answer:

You can get the indices of the upper triangle in your matrix using:

indices = np.triu_indices_from(A)

indices
Out[1]: 
(array([0, 0, 0, 1, 1, 2], dtype=int64),
 array([0, 1, 2, 1, 2, 2], dtype=int64))

This will include the diagonal indices, to exclude them you can offset the diagonal by 1:

indices_with_offset = np.triu_indices_from(A, k=1)

indices_with_offset
Out[2]: 
(array([0, 0, 1], dtype=int64), 
 array([1, 2, 2], dtype=int64))

Now use these with your matrix as a mask

A[indices_with_offset]

Out[3]: 
array([2, 3, 6])

See docs here

Answer 2

One approach with masking -

def upper_tri_masking(A):
    m = A.shape[0]
    r = np.arange(m)
    mask = r[:,None] < r
    return A[mask]

Another with np.triu_indices -

def upper_tri_indexing(A):
    m = A.shape[0]
    r,c = np.triu_indices(m,1)
    return A[r,c]

Sample run -

In [403]: A
Out[403]: 
array([[79, 17, 79, 58, 14],
       [87, 63, 89, 26, 31],
       [69, 34, 90, 24, 96],
       [59, 60, 80, 52, 46],
       [75, 80, 11, 61, 47]])

In [404]: upper_tri_masking(A)
Out[404]: array([17, 79, 58, 14, 89, 26, 31, 24, 96, 46])

Runtime test -

In [415]: A = np.random.randint(0,9,(5000,5000))

In [416]: %timeit upper_tri_masking(A)
10 loops, best of 3: 64.2 ms per loop

In [417]: %timeit upper_tri_indexing(A)
1 loop, best of 3: 252 ms per loop