tad.py 文件源码

def kth_diag_indices(n, k):
    """
    Return indices of bins k steps away from the diagonal.
    (from http://stackoverflow.com/questions/10925671/numpy-k-th-diagonal-indices)
    """

    rows, cols = np.diag_indices(n)
    if k < 0:
        return rows[:k], cols[-k:]
    elif k > 0:
        return rows[k:], cols[:-k]
    else:
        return rows, cols