[백준 17298][파이썬] 오큰수

https://www.acmicpc.net/problem/17298

17298번: 오큰수

우선 문제를 봤을 때 가장 먼저 생각나는 알고리즘은 수열의 모든 값을 조사하는데 각각의 수열마다 자신 다음 값부터 끝까지 자신보다 큰 값이 나올 때까지 찾는 방법이다.

 

9 8 7 6 ... 과 같이 내림차순으로 정렬된 수열이 최악의 경우인데 이럴 경우 1부터 n까지 더한 횟수를 반복해야 한다.

 

즉 시간 복잡도가 O(N^2)이 된다.

 

n의 최대값이 1,000,000이고 파이썬은 1초에 1억번 정도의 연산을 할 수 있다는 점을 고려하면 이 방법은 시간 초과가 발생한다.

 

이 시간을 줄이기 위해서는 최소 힙을 사용해야 한다.

 

예제 입력 2번을 보자.

 

우선 수열의 첫 번째 값부터 조사를 시작한다.

 

1.

첫 번째 값은 무조건 최소 힙에 넣는다. (편의상 힙의 모양을 오름차순으로 정렬된 리스트로 표현하겠다.)

 

heap : [ 9 ] 

 

2.

최소 힙에서 가장 작은 값(temp라고 표현하겠다.)을 꺼내서 두 번째 값 (arr[2] 라고 표현하겠다.)과 비교한다.

 

heap : [ ]

temp = 9, arr[2] = 5

 

temp가 arr[2]보다 크다. 이럴 경우 다시 heap에 temp를 집어넣는다.

 

arr[2] 값도 heap에 넣는다.

 

heap : [ 5, 9 ]

 

3.

최소 힙에서 가장 작은 값(temp)을 꺼내서 세 번째 값  (arr[3]) 과 비교한다.

 

heap : [ 9 ]

temp : 5, arr[3] = 4

 

temp가 arr[3]보다 크다. 이럴 경우 다시 heap에 temp를 집어넣는다.

 

arr[3] 값도 heap에 넣는다.

 

heap : [ 4, 5, 9 ]

 

4.

최소 힙에서 가장 작은 값(temp)을 꺼내서 네 번째 값  (arr[4]) 과 비교한다.

heap : [ 5, 9 ]

temp = 4, arr[4] = 8

 

드디어 temp가 arr[4]보다 작다. 이럴 경우 temp에 해당하는 곳의 오큰수는 arr[4]가 된다.

 

temp가 arr[4]보다 작을 경우 heap에서 한 번 더 가장 작은 값을 꺼낸다.

 

오큰수가 arr[4]인 값이 있을 수도 있기 때문이다.

 

heap : [ 9 ]

temp = 5, arr[4] = 8

 

또 temp가 arr[4]보다 작다. temp에 해당하는 곳의 오큰수는 arr[4]가 된다.

 

한 번 더 꺼낸다.

 

heap : [  ]

temp = 9, arr[4] = 8

 

temp가 arr[3]보다 크다. 이럴 경우 다시 heap에 temp를 집어넣는다.

 

heap : [ 9 ]

 

이제 모든 수열의 값을 조사하였다.

 

아직도 heap에 남아있는 값이 있다면 오큰수가 존재하지 않는 것이다.

 

<전체 코드>

import sys
from collections import deque
import heapq
input=sys.stdin.readline
n=int(input())
arr=list(map(int, input().split()))
heap=[(arr[0], 0)] #첫 번째는 값 두 번째는 인덱스
res=[-1]*n #각각의 인덱스에 대한 오큰수 값
for i in range(1, n):
    while True:
        temp=heapq.heappop(heap) #가장 작은 값 꺼냄
        if temp[0]<arr[i]: #꺼낸 값이 현재 위치보다 작으면
            res[temp[1]]=arr[i] #오큰수 저장
        else:
            heapq.heappush(heap, temp) #꺼낸 값이 크면 다시 최소 힙에 넣음
            break
        if len(heap)==0:
            break
    
    heapq.heappush(heap, (arr[i], i)) #최소 힙에 넣기
for i in res:
    print(i, end=' ')