NumPy k-th diagonal indices

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

Here's a way:

1. Create index value arrays.
2. Get the daigonal index values you want.
3. Thats it! :)

Like this:

>>> import numpy as np>>> rows, cols = np.indices((3,3))>>> row_vals = np.diag(rows, k=-1)>>> col_vals = np.diag(cols, k=-1)>>> z = np.zeros((3,3))>>> z[row_vals, col_vals]=1>>> zarray([[ 0.,  0.,  0.],       [ 1.,  0.,  0.],       [ 0.,  1.,  0.]])

The indices of the k'th diagonal of a can be computed with

def kth_diag_indices(a, k):    rowidx, colidx = np.diag_indices_from(a)    colidx = colidx.copy()  # rowidx and colidx share the same buffer    if k > 0:        colidx += k    else:        rowidx -= k    k = np.abs(k)    return rowidx[:-k], colidx[:-k]

Demo:

>>> aarray([[ 0,  1,  2,  3,  4],       [ 5,  6,  7,  8,  9],       [10, 11, 12, 13, 14],       [15, 16, 17, 18, 19],       [20, 21, 22, 23, 24]])>>> a[kth_diag_indices(a, 1)]array([ 1,  7, 13, 19])>>> a[kth_diag_indices(a, 2)]array([ 2,  8, 14])>>> a[kth_diag_indices(a, -1)]array([ 5, 11, 17, 23])