Brandon - 6 months ago 28
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
``````

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])
``````
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download