Brandon Brandon - 26 days ago 9
Python Question

Write a loop (or a nested loop) that extracts the diagonal elements 2,3,4,6,7,8 and prints them one by one

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 Source

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])