Upper diagonal of a 2D numpy array

This seems simple (and is trivial to write a three-line loop for), b开发者_如何学Gout how can I use numpy slicing make a list of the index locations of the upper diagonal of a numpy array? I.e.

Given a 4x4 array, I'd like the index locations at the X's:

[ X X X X ]
[ 0 X X X ]
[ 0 0 X X ]
[ 0 0 0 X ]

Giving:

[ (0,0), (0,1), (0,2), (0,3), (1,1), (1,2), (1,3), (2,2), (2,3), (3,3) ]

carnieri beat me to the numpy.triu_indices answer, but there is also numpy.triu_indices_from which takes an array as input rather than the dimensions.

Though the format of index locations is different, it seems like you want the function numpy.triu_indices.

scipy.linalg.triu and nonzero

If you're running Ubuntu and don't want to upgrade Numpy just for this, you can use the following function:

from itertools import chain
triu_indices = lambda x, y=0: zip(*list(chain(*[[(i, j) for j in range(i + y, x)] for i in range(x - y)])))

Example:

In [26]: triu_indices = lambda x, y=0: zip(*list(chain(*[[(i, j) for j in range(i + y, x)] for i in range(x - y)])))

In [27]: triu_indices(4)
Out[27]: [(0, 0, 0, 0, 1, 1, 1, 2, 2, 3), (0, 1, 2, 3, 1, 2, 3, 2, 3, 3)]

In [28]: zip(*triu_indices(4))
Out[28]: 
[(0, 0),
 (0, 1),
 (0, 2),
 (0, 3),
 (1, 1),
 (1, 2),
 (1, 3),
 (2, 2),
 (2, 3),
 (3, 3)]