[프로그래머스] n^2 배열 자르기 - python

n^2 배열 자르기

문제

정수 n, left, right가 주어집니다. 다음 과정을 거쳐서 1차원 배열을 만들고자 합니다.

1. n행 n열 크기의 비어있는 2차원 배열을 만듭니다.

2. i = 1, 2, 3, ..., n

   에 대해서, 다음 과정을 반복합니다.
   • 1행 1열부터 i행 i열까지의 영역 내의 모든 빈 칸을 숫자 i로 채웁니다.
3. 1행, 2행, ..., n행을 잘라내어 모두 이어붙인 새로운 1차원 배열을 만듭니다.
4. 새로운 1차원 배열을 arr이라 할 때, arr[left], arr[left+1], ..., arr[right]만 남기고 나머지는 지웁니다.

정수 n, left, right가 매개변수로 주어집니다. 주어진 과정대로 만들어진 1차원 배열을 return 하도록 solution 함수를 완성해주세요.

제한사항

• 1 ≤ n ≤ 10^7
• 0 ≤ left ≤ right < n^2
• right - left < 10^5

풀이

처음에는 n * n 배열 전체를 구했는데 시간 초과가 났다. O(n^2)이하로 만들 수 있는 방법이 떠오르지 않아서 다른 사람 풀이를 참고했다.

n = 3일 때를 생각해보면

 

1 2 3 (0,0) (0,1) (0,2)

2 2 3 (1,0)(1,1)(1,2)

3 3 3 (2,0)(2,1)(2,2)

 

arr은 아래와 같은 형태가 된다.

 

   0       1        2      3      4     5       6       7      8

(0,0) (0,1) (0,2) (1,0) (1,1) (1,2) (2,0) (2,1) (2,2)

   1       2       3      2       2      3      3       3      3

 

arr[0] = max(0//3 , 0%3) + 1 = 1

arr[1] = max(1//3 , 1%3) + 1 = 2

arr[2] = max(2 // 3, 2 % 3) + 1 = 3

.

.

.

 

일반화해 보면 arr[i] = max(i // n, i % n) + 1 임을 알 수 있고 따라서 left, right 사이의 값만 계산해서 넣어주면 된다.

테스트케이스에 n =1일 경우는 없는 것 같지만 그냥 예외 처리 해줬다.

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