# numpy.diag_indices

numpy.diag_indices(n, ndim=2)

  Return the indices to access the main diagonal of an array.

  This returns a tuple of indices that can be used to access the main diagonal of an array a with `a.ndim >= 2` dimensions and shape (n, n, …, n). For `a.ndim = 2` this is the usual diagonal, for `a.ndim > 2` this is the set of indices to access `a[i, i, ..., i]` for `i = [0..n-1]`.

  Parameters
    n int

      The size, along each dimension, of the arrays for which the returned indices can be used.

    ndim int, optional

      The number of dimensions.

  See also

  `diag_indices_from`

  Notes

  New in version 1.4.0.

  Examples

  Create a set of indices to access the diagonal of a (4, 4) array:

      >>> di = np.diag_indices(4)
      >>> di
      (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
      >>> a = np.arange(16).reshape(4, 4)
      >>> a
      array([[ 0,  1,  2,  3],
             [ 4,  5,  6,  7],
             [ 8,  9, 10, 11],
             [12, 13, 14, 15]])
      >>> a[di] = 100
      >>> a
      array([[100,   1,   2,   3],
             [  4, 100,   6,   7],
             [  8,   9, 100,  11],
             [ 12,  13,  14, 100]])

  Now, we create indices to manipulate a 3-D array:

      >>> d3 = np.diag_indices(2, 3)
      >>> d3
      (array([0, 1]), array([0, 1]), array([0, 1]))

  And use it to set the diagonal of an array of zeros to 1:

      >>> a = np.zeros((2, 2, 2), dtype=int)
      >>> a[d3] = 1
      >>> a
      array([[[1, 0],
              [0, 0]],
             [[0, 0],
              [0, 1]]])

previous

numpy.unravel_index

next

numpy.diag_indices_from