NumPy k-th diagonal indices

Question

I\'d like to do arithmetics with k-th diagonal of a numpy.array. I need those indices. For example, something like:

>>> a = numpy.eye(2)
>>>         


        

Answer 1

Use numpy.diag(v, k=0)

Where k sets the diagonal location from center.

ie. {k=0: "default center", k=(-1): "1 row to the left of center", k=1: "1 row to the right of center}

Then perform the arithmetic as you would normally expect.

Check out the docs here: np.diag().

Examples:

In [3]: np.diag(np.arange(6), k=0)
Out[3]: 
array([[0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0],
       [0, 0, 0, 3, 0, 0],
       [0, 0, 0, 0, 4, 0],
       [0, 0, 0, 0, 0, 5]])

In [4]: np.diag(np.arange(6), k=1)
Out[4]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 1, 0, 0, 0, 0],
       [0, 0, 0, 2, 0, 0, 0],
       [0, 0, 0, 0, 3, 0, 0],
       [0, 0, 0, 0, 0, 4, 0],
       [0, 0, 0, 0, 0, 0, 5],
       [0, 0, 0, 0, 0, 0, 0]])

In [5]: np.diag(np.arange(6), k=-1)
Out[5]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0, 0],
       [0, 0, 0, 3, 0, 0, 0],
       [0, 0, 0, 0, 4, 0, 0],
       [0, 0, 0, 0, 0, 5, 0]])