[Python] 특별한 이차원 배열 1

문제 설명

문제 설명 정수 n이 매개변수로 주어질 때, 다음과 같은 n × n 크기의 이차원 배열 arr를 return 하는 solution 함수를 작성해 주세요.

예시 ) 

n          result

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

 

# 단위행렬 만들기
def solution(n):
    arr = []
    for i in range(n):
        row = []
        for j in range(n):
            # i와 j가 같은 인덱스일때만 1을 삽입한다. 단위행렬의 경우 [i][j] 가 1이다. 
            if i == j:
                row.append(1)
            else:
                row.append(0)
        # arr 안에 반복문이 끝날때마다 row 를 삽입한다. 
        arr.append(row)
    return arr