Extract diagonals of each 2D square sub-array from a larger 2D array

Question

This is my first programming class and I'm very excited to learn Python for Data Science. I cannot figure out how to write a loop that returns all diagonal numbers in the matrix. Below is the code, how close or far off am I? Thank you!

import numpy as np
cols = 0

matrixA = np.array([[2,0,0], [0,3,0], [0,0,4], [6,0,0], [0,7,0], [0,0,8]])

for rows in range(6):
    if rows == cols:
        print(matrixA[rows, cols])
    cols = cols + 1

Answer 1

Your current solution does not work because it does not take into account the fact that matrixA is not square. You will have to take care that your indices do not run out of bounds. Running it gives:

IndexError: index 3 is out of bounds for axis 1 with size 3

This is because the maximum value that cols is allowed to take here is 2.

---

As an alternative, you could use np.diag:

print(x)
array([[2, 0, 0],
       [0, 3, 0],
       [0, 0, 4],
       [6, 0, 0],
       [0, 7, 0],
       [0, 0, 8]])

res = np.array([np.diag(x, -offset) for offset in range(0, *x.shape)])

print(res)
array([[2, 3, 4],
       [6, 7, 8]])

If you want a 1D result, call np.ravel:

print(res.ravel())
array([2, 3, 4, 6, 7, 8])