andrej-voli-sladoled.py

import numpy as np

a = np.zeros((3, 3))
print(a)

# row i col indexi
ri, ci = np.diag_indices_from(a)
# na row i col indexe ubacimo odgovarujuce elemente tim redom
a[ri, ci] = (1, 2, 3)
print(a)