[Python/파이썬] 프로그래머스 n^2 배열 자르기

프로그래머스

코드

• 처음 시도했던 코드 (시간 초과!)

def solution(n, left, right):
    answer = []
    dim2 = [[0 for _ in range(n)] for _ in range(n)]
    for i in range(n):
        for j in range(i, n):
            dim2[i][j] = j+1
            dim2[j][i] = j+1

    for row in dim2:
        answer.extend(row)
    return answer[left:right+1]

문제에 나와 있는대로 2차원 배열을 만든 뒤 그걸 그대로 1차원 배열로 만들어서 잘라주는 코드이다. 당연히 될거란 기대도 안했는데 그냥 해봤다. 역시나 시간 초과로 통과하지 못한다. 그래서 다음과 같이 풀이하였다.

def solution(n, left, right):
    answer = []
    for i in range(left, right+1):
        row, col = i // n, i % n
        if row < col:
            answer.append(col+1)
        else:
            answer.append(row+1)
    return answer

문제의 2차원 배열을 잘 살펴보면 열보다 행의 인덱스가 작을 때 그 칸에는 행의 인덱스가 저장되고 그렇지 않다면 열의 인덱스가 저장되는 것을 알 수 있다. 그래서 left부터 right까지 행과 열의 인덱스를 구하여 두 인덱스값을 비교해보며 알맞는 값을 answer 배열에 추가해준다. 이때 주의할 점은 문제는 0행 0열에 0이 아닌 1부터 시작하여 저장하므로 answer 배열에 값을 추가할 때는 1을 더한 값을 넣어주어야 한다.