[Level 2] 뒤에 있는 큰 수 찾기(Python)

프로그래머스

문제

정수로 이루어진 배열 numbers
모든 원소에 대해 뒷 큰수들을 차례대로 담은 배열을 return
(단, 뒷 큰수가 존재하지 않는 원소는 -1을 담는다)

<제한 사항>

- 4 ≤ numbers의 길이 ≤ 1,000,000
    - 1 ≤ numbers[i] ≤ 1,000,000

 

풀이 과정

풀이 시간: 25분

알고리즘: 자료구조(스택) or 우선순위큐(힙)

numbers의 길이가 최대 1,000,000이기 때문에 최악의 경우 O(NlogN)으로 풀이해야 한다.

 

1. stack을 이용해 O(N)으로 풀이할 경우
맨 뒤의 수는 항상 뒷큰수를 가질 수 없기 때문에 -1로 고정하고, 뒤에서부터 역순으로 탐색한다.
각 원소는 한 번 추가되고 한 번 제거되므로 총 제거 작업의 횟수는 최대 'n'번

 

2. heap을 이용해 O(NlogN)으로 풀이할 경우

numbers의 배열을 처음부터 끝까지 순회한다 = O(N)

현재 원소 numbers[i]가 힙의 최솟값보다 크면 힙에서 최솟값을 제거(=O(logN))하고, 해당 인덱스에 현재 값을 answer에 저장한다.

힙에 numbers[i]와 i를 최소힙으로 push한다. = O(logN)

 

스택을 이용한 풀이

def solution(numbers: list) -> list:
    answer = [-1] * len(numbers)
    stack = [numbers[-1]]
    for i in range(len(numbers)-2, -1, -1):
        # 스택에 들어온 순서대로 탐색, numbers[i]보다 큰 원소를 발견할 때까지
        while stack and stack[-1] <= numbers[i]:
            stack.pop() # 원소 제거, 어차피 다음 순회할 원소 입장에서도 필요 없는 수이기 때문
        if stack:
            answer[i] = stack[-1]
        stack.append(numbers[i])
    return answer

 

 

힙을 이용한 풀이(그리디 알고리즘)

import heapq


def solution(numbers):
    answer = [-1] * len(numbers)
    heap = []
    for i in range(len(numbers)):
        while heap and heap[0][0] < numbers[i]:
            value, idx = heapq.heappop(heap)
            answer[idx] = numbers[i]
        heapq.heappush(heap, (numbers[i], i)) # numbers의 수를 최소힙으로 push
    return answer

 

백준의 오큰수와 거의 똑같은 문제로 스택을 이용해 풀 수 있는 문제였다. 

heap으로도 풀이가 가능한데 heap의 경우 O(NlogN)으로 런타임이 약간의 느리지만 좀 더 깔끔한 풀이가 가능하다.