[BOJ] 2075 N번째 큰 수

문제

N×N의 표에 수 N2개 채워져 있다. 채워진 수에는 한 가지 특징이 있는데, 모든 수는 자신의 한 칸 위에 있는 수보다 크다는 것이다. N=5일 때의 예를 보자.

12	7	9	15	5
13	8	11	19	6
21	10	26	31	16
48	14	28	35	25
52	20	32	41	49

이러한 표가 주어졌을 때, N번째 큰 수를 찾는 프로그램을 작성하시오. 표에 채워진 수는 모두 다르다.

입력

첫째 줄에 N(1 ≤ N ≤ 1,500)이 주어진다. 다음 N개의 줄에는 각 줄마다 N개의 수가 주어진다. 표에 적힌 수는 -10억보다 크거나 같고, 10억보다 작거나 같은 정수이다.

출력

첫째 줄에 N번째 큰 수를 출력한다.

풀이

보기엔 굉장히 쉬운 문제였으나 처음엔 굉장히 복잡하게 생각했다. 일단 우선순위큐를 써야하는 것은 확실했기에(또한 나는 처음에 모든 요소를 스캔하면 시간초과가 날 것이라고 생각했다... 1500*1500을 계산 안한 자의 최후다)나는 배열의 뒤부터 시작해서 배열이 비어있으니 맨 첫 줄을 heapq.heappush()로 넣어주고(이 때 값만 넣어주는 것이 아니라 그 요소를 타고 올라가기 위해 x좌표 y좌표 또한 넣어주었다), heapq.heappop()으로 얻은 좌표로 위로 따라 올라가면서 만약 heap안의 가장 큰 값보다 가장 작은 값을 가지고 있으면 heap에 넣어주고 탈출하는 식으로 풀었다. 정답은 나오는 코드였지만 보다시피 메모리가 극도로 적게 주어지는 문제였기에(12MBㄷㄷ) 당연히 메모리 초과가 났다. 이제 최대 힙이 아니라 최소 힙을 이용해서 풀기로 하고, 주어지는 메모리 양을 보아하니 힙 안에는 N개의 요소만 들어있어야 할 것 같았다. 먼저 힙이 비어있으면 한 줄을 heappush로 넣어주고, 비어있지 않다면 요소를 비교하여 만약 heap의 가장 작은 요소보다 크다면 heappop을 해주고 다시 넣어주었다. 이런식으로 하다보면 결국엔 heap에는 N개의 요소만 남고, 그 중에 가장 작은 값을 출력하면 된다.

메모리초과 코드

import sys
import heapq
input = lambda: sys.stdin.readline().rstrip()
n = int(input())
cnt = 0
matrix = [list(map(int, input().split()))for _ in range(n)]
arr = []
for i in range(n):
  heapq.heappush(arr, (-matrix[n-1][i], n-1, i))
while True:
  if cnt == n-1:
    break
  a = heapq.heappop(arr)
  cnt += 1
  for i in range(a[1]-1, -1, -1):
    if matrix[i][a[2]] > arr[0][0]:
      heapq.heappush(arr, (-matrix[i][a[2]], i, a[2]))
      break
    cnt += 1
    if cnt == n:
      heapq.heappush(arr, (-matrix[i][a[2]], i, a[2]))
      break
print(-arr[0][0])

정답 코드

import sys
import heapq
input = lambda: sys.stdin.readline().rstrip()
n = int(input())
arr = []
for _ in range(n):
  numbers = list(map(int, input().split()))
  if not arr:
    for num in numbers:
      heapq.heappush(arr, num)
  else:
    for num in numbers:
      if num > arr[0]:
        heapq.heappop(arr)
        heapq.heappush(arr, num)
print(arr[0])