[백준] 2075번 : N번째 수 파이썬

시간제한	메모리제한	정답 비율
1 초	12 MB	39.455%

문제

N×N의 표에 수 N^2개 채워져 있다. 채워진 수에는 한 가지 특징이 있는데, 모든 수는 자신의 한 칸 위에 있는 수보다 크다는 것이다. 아래와 같이 표가 주어졌을 때, N번째 큰 수를 찾는 프로그램을 작성하시오. 표에 채워진 수는 모두 다르다.

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

출력

첫째 줄에 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

예제 출력

35

해결 아이디어

N의 범위가 주어지기를 N(1 ≤ N ≤ 1,500)이고, 각줄에 주어진 수의 범위가 -10억에서 10억까지이다. 메모리 제한이 12mb이므로, 2차원 배열이나 튜플을 이용해 최대힙을 구현할 경우 시간 초과가 발생할 것이다.  

파이썬 heapq 모듈을 통해 힙을 만들면, 최소힙이 만들어지고 최소힙 구조에서 조건문을 사용하여 힙 원소의 push, pop을 제어해 최소한의 메모리 사용을 하는 것이 목표이다. 

이제, 문제를 해결하기 위한 알고리즘 아이디어를 짜보자. 문제에서 주어진 예제와 같이 n을 5라고 가정하고, 첫번째 입력으로 "12 7 9 15 5"가 주어졌다고 해보자.

최초로 입력받은 숫자들

 우선, 첫번째 입력받은 숫자들은 대조할 대상이 없기 때문에 그대로 heapify하여 heap에 저장한다. 

최초의 heapify

그러면 이렇게 최소힙이 만들어지는데, 이중 n번째 큰 수라고 하면 '5'가 답이 될 것이다. 하지만 우리는 n*n사이즈의 표에서 n번째 큰 수를 찾아야하고, 다음 입력이 주어지니까 다음 입력을 받아보자. 

heap은 현재 오름차순으로 정렬되어있는 상태이고, 입력받은 숫자들은 정렬되지 않은 상태이다. 

heappush()를 이용하면, 우선 순위에 따른 정렬이 일어나므로 입력받은 숫자들을 바로 heap의 원소들과 대조하여 보자. 

입력의 첫번째 원소인 '13'을 heap의 0번째 원소와 대조해보자. 만일 input의 0번째 숫자인 '13'이 heap[0]인 '5'보다 크다면 heap에 '13'을 heappush하고, heappop()을 수행한다. 그럼 아래 이미지와 같이 heap에는 [7, 9, 12, 13, 15]가 남는다. 이 상태는 입력받은 모든 숫자들 중에서 1~5번째 큰 수가 힙에 남아있는 상태이다. 입력에 남아있는 다른 숫자들도 계속해서 대조 연산을 실행해보자. 

마지막 입력된 숫자 '6'처럼 heap의 0번째 원소보다 작은 수라면, 큰 수가 아니므로 heap에 push할 필요도 없다. 

이제 3번째 숫자들을 입력받자. 이 상태가 되었다면, 2번째 입력 숫자들 중에서 1~n번째 큰 수들을 골라낸 상태이다. 

현재까지 입력받았던 숫자들은 [12, 7, 9, 15, 5, 13, 8, 11, 19, 6] 이고 이들을 오름차순으로 나타내면 [5, 6, 7, 8, 9, 11, 12, 13, 15, 19]이므로 [11, 12, 13, 15, 19]는 이전 입력까지 상위 5개의 숫자만이 heap에 남아있는 상태이다. 이제 앞서 진행했던 방식과 같이 남은 입력들을 처리하면, 마지막에 heap[0]에 남아있는 원소가 모든 입력된 수들 중에서 n번째로 큰 수가 된다.

3번째로 입력받은 숫자들과 이전까지 입력받은 수들 중 1~5번째로 큰 수들이 남아있는 heap의 상태

모든 과정이 끝난 후 heap의 상태

위 이미지는 모든 입력에 대해 원소 대소 크기 비교 및 heappush, heappop 연산이 끝난 후의 상태이다. 문제에서 원하는 출력이 n번째 큰 수 였으므로 이 상태에서 heap[0]을 출력하면 답을 얻을 수 있다. 

 

import heapq as hq # 힙큐를 hq로 쓰기
import sys
input = sys.stdin.readline

n = int(input())
heap = []

init_numbers = list(map(int,input().split())) 
for num in init_numbers:
            hq.heappush(heap,num) #list 객체를 풀어서 원소를 heap에 넣어준다.

for i in range(n-1):  
    numbers = list(map(int,input().split())) 
        # heap의 0번째 원소는 해당 입력까지, 모든 수 들 중 n번째 큰 수
        # 입력 한줄 마다 힙 원소를 뒤바꾼다.
    for num in numbers:
        if heap[0] < num :
            hq.heappush(heap,num)
            hq.heappop(heap)

print(heap[0])