identity = [[1 if i == j else 0 for j in range(3)] for i in range(3)]
print(identity)


[[1, 0, 0], [0, 1, 0], [0, 0, 1]]