NumPy k-th diagonal indices

Question

I\'d like to do arithmetics with k-th diagonal of a numpy.array. I need those indices. For example, something like:

>>> a = numpy.eye(2)
>>>         


        

Answer 1

A bit late, but this version also works for k = 0 (and does not alter the arrays, so does not need to make a copy).

def kth_diag_indices(a, k):
    rows, cols = np.diag_indices_from(a)
    if k < 0:
        return rows[-k:], cols[:k]
    elif k > 0:
        return rows[:-k], cols[k:]
    else:
        return rows, cols