[백준 2075] N번째 큰 수

 

 

블로그 개설 후 첫 번째 문제다. 앞으로 실버 이상의 문제들을 올릴 생각이고 열심히 해보겠습니다.

이 문제를 처음 풀었을 때 단순히 메모리 생각을 안 하고 풀었지만 메모리 초과가 떠서 보니 메모리 제한 12MB였다,,,

import sys
input = sys.stdin.readline

lst = []
N = int(input())
for i in range(N):
    for j in range(N):
        lst.append(input().rstrip().split())

lst = list(map(int,lst))
lst.sort()
print(lst[-N])

일단 백준을 풀면서 느낀 건 입력을 받을 때는 그냥 import sys 하고 sys.stdin.readline을 이용하는 것이 마음이 편할 것 같다.

괜히 input으로 시간초과 날빠에는 sys.stdin.readline으로 모두 입력을 받아 버리자.

 

처음 생각한 방법은 input(). rstrip(). split()으로 개행문자를 제거하고 모든 수를 배열에 때려 박은 다음에 정렬 후 뒤에서 N번째 수를 출력하려 했는데 메모리 초과가 나버렸다.

 

이 문제는 heapq를 사용하는 문제인데 파이썬에서 힙을 처음 사용해 봐서 이번 기회에 배우고 정리를 해보려 한다.

파이썬에서는 import heapq를 통해 힙 자료구조를 사용할 수 있다.

• heapq.heappush(heap, item) : item을 heap에 추가
• heapq.heappop(heap) : heap에서 가장 작은 원소를 pop & 리턴. 비어 있는 경우 IndexError가 호출됨. 
• heapq.heapify(x) : 리스트 x를 즉각적으로 heap으로 변환함 (in linear time, O(N) )

-힙 생성 & 힙 원소 추가

import heapq

heap = []
heapq.heappush(heap, 50)
heapq.heappush(heap, 10)
heapq.heappush(heap, 20)

print(heap)

 

---> 출력 : [10, 50, 20]

 

최소 힙이 생성되므로 첫 번째 원소가 10인 것을 확인할 수 있다.

만약 최소 힙이 아니라면 heapify(리스트)를 통해 최소 힙으로 바꿀 수 있다.

 

-힙 원소 삭제

heapq.heappop(리스트)

 

다음은 최대힙에 대해 알아보겠다.

heap_items = [1,3,5,7,9]

max_heap = []
for item in heap_items:
  heapq.heappush(max_heap, (-item, item))

print(max_heap)

 

---> 출력 : [(-9, 9), (-7, 7), (-3, 3), (-1, 1), (-5, 5)]

힙에 push 할 때 (-item, item)으로 튜플형태로 추가하면 최대힙을 만들 수 있다.

실제 원소의 값을 알고 싶을 때는 heapq.heappop(max_heap)[1]을 이용하면 된다.

 

• 최종 제출 코드

import sys
import heapq
input = sys.stdin.readline

n = int(input())
heap = []
for _ in range(n):
    for i in list(map(int,input().split())):
        heapq.heappush(heap, i) # 최소힙
        if len(heap) > n:
            heapq.heappop(heap)
            
print(heapq.heappop(heap))

 

먼저 숫자를 입력받으면서 heap에 모두 heappush 해준다. 이때 heap의 길이가 n보다 크다면 heapq.heappop을 진행한다.

이러면 힙은 길이를 n개를 유지하면서 현재 입력받은 수 중 가장 큰 n개만 남게 된다.

마지막으로 heapq.heappop(heap)을 출력하면 남은 n개 중 가장 작은 수가 출력되는데 이 수가 곧 N번째 큰 수이다.