fast way to get the indices of a lower triangular matrix as 1 dimensional list in python

Question

Given the number of rows (or columns) , n, of a square matrix, I am trying to get the index pairs of the lower triangular matrix in a 1 dimensional list. So far I thought of the following solution:

def getLowerTriangularIndices(n):
    inds=[];         
    for i in range(1,n):
        for j in range(i):
            inds.append((i,j))
    return inds;

Considering the two for loops, it would be far better to have a more efficient way of calculating this maybe using numpy. Does anyone have a suggestion?

Answer 1

Numpy has a method for that...

import numpy as np

# create your matrix. If it's not yet a numpy array, make it one
ar = np.array(matrix)
indices = np.tril_indices_from(ar)

This returns a tuple of two arrays. If you want to have them as lists, you could do

indices = [list(x) for x in np.tril_indices_from(ar)]

You actually do not need to have an array to get the indices, there is also np.tril_indices, which takes the shape as arguments.

So your function would read:

def getLowerTriangularIndices(n):
    return [list(x) for x in np.tril_indices(n)]

or if you want a list of tuples instead:

def getLowerTriangularIndices(n):
    return zip(*np.tril_indices(n)]