NumPy第k个对角线索引

NumPy第k个对角线索引

有点晚了，但是此版本也适用

k = 0

（并且不会更改数组，因此不需要进行复制）。

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