+2 votes
in Programming Languages by (17.9k points)
How can I extract the diagonal elements of a 2D array?

E.g.

[[1,2,3],[4,5,6],[7,8,9]]

Ans: [1,5,9]

1 Answer

+1 vote
by (28.5k points)

You can use the diag() function of Numpy to extract the diagonal elements.

Here is an example:

>>> import numpy as np
>>> a=np.random.randn(3,3)
>>> a
array([[-0.17423025,  0.42428257, -0.420369  ],
       [ 0.66431584,  1.25155595,  1.84738162],
       [-0.58160319,  0.67047704,  1.0755773 ]])
>>> np.diag(a)
array([-0.17423025,  1.25155595,  1.0755773 ])
>>>
 


...