Search code examples
pythonnumpyindices

Generating indices using np.triu_indices

I would like call a function on each entry of an upper-triangular matrix. In order to avoid a messy nest of for loops, I am using the numpy.triu_indices function. My function works but I would like to know if there is a cleaner way to implement the indexing. 

import numpy as np
def weighted_edges(adjmat):
    indices = np.triu_indices(len(adjmat))
    return ((x, y, adjmat[x,y]) for (x,y) in zip(indices[0], indices[1]))

I suspect that there is a way to implement this without needing to reference indices[i] in the zip call. Is there indeed a way to do so?

Solution

  • If you have an N x N matrix from which you want the upper triangular values, just do

    import numpy as np
N = 5
x = np.arange(N**2).reshape(N, N)
upper = x[np.triu_indices(N, 0)]

    If you want the triangular values offset from the main diagonal by k columns, then do

    upper = x[np.triu_indices(N, k)]