matrix.py 文件源码

def kth_diag_indices(array, diag_k):
    """Return a tuple of indices for retrieving the k'th diagonal
    of matrix a.

    :param array: Input matrix.
    :type array: :class:`numpy array <numpy.ndarray>`
    :param int diag_k: Diagonal to index. 0 is the centre, 1 is the first \
            diagonal below the centre, -1 is the first diagonal above \
            the centre.

    >>> my_matrix = np.array([[ 0, -1, -2, -3],
    ...                       [ 1,  0, -1, -2],
    ...                       [ 2,  1,  0, -1],
    ...                       [ 3,  2,  1,  0]])
    >>> matrix.kth_diag_indices(my_matrix, 1)
    (array([1, 2, 3]), array([0, 1, 2]))
    >>> my_matrix[
    ...    matrix.kth_diag_indices(my_matrix, 1)
    ...  ]
    array([1, 1, 1])
    >>> my_matrix[
    ...    matrix.kth_diag_indices(my_matrix, -2)
    ...  ]
    array([-2, -2])
    """

    rows, cols = np.diag_indices_from(array)
    if diag_k < 0:
        return rows[:diag_k], cols[-diag_k:]

    if diag_k > 0:
        return rows[diag_k:], cols[:-diag_k]

    return rows, cols