def solution(n):
matrix = [[1 if i == j else 0 for j in range(n)] for i in range(n)]
return matrix
방법 1:
def solution(n):
answer = [[0] * n for i in range(n)]
for i in range(n):
answer[i][i] = 1
return answer