[프로그래머스] 뒤에있는 큰 수 찾기 (Python)
문제설명
정수로 이루어진 배열 numbers가 있습니다.
배열 의 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 합니다.
정수 배열 numbers가 매개변수로 주어질 때,
모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 return 하도록 solution 함수를 완성해주세요.
단, 뒷 큰수가 존재하지 않는 원소는 -1을 담습니다.
제한사항
4 ≤ numbers의 길이 ≤ 1,000,000
1 ≤ numbers[i] ≤ 1,000,000
문제풀이
일단 이 문제는 프로그래밍 개념 중 하나인 자료구조를 공부하는 사람들에게
매우 도움이 될 문제이다.
문제 풀이를 나아가기 전, 해당 문제의 시간복잡도를 보고 어떠한 자료구조를 활용할지 대강 생각해보자
numbers의 길이가 최대 100만이다.
따라서, N^2 알고리즘을 사용하게 된다면, 반드시 시간초과가 발생할 것이다.
해당 문제에서 제한사항을 제대로 보지 못했다면, 시간초과가 발생했을 것이다.
아래와 같은 코드와 함께 제출하면서 말이다.
def solution(numbers):
answer = []
for i, num in enumerate(numbers):
found = False
for n in numbers[i+1:]:
if n > num:
answer.append(n)
found = True
break
if not found:
answer.append(-1)
return answer
본론으로 들어가서 문제를 해결하는 방법에 대해 알아보자
일단, 시간 안에 문제를 해결하기 위해선 최소 NlogN의 알고리즘을 작성해야한다.
(적어도 위와 같은 이중 반복문을 사용할 수 없다는 점을 고려해야 한다.)
그렇다면 문제를 해결하기 위해선 어떻게 해야 할까?
1. heap 구조를 활용한 해결방법
가장 쉬운 아이디어로는 heap 자료구조를 활용하는 방법이 있다.
heap 자료구조는 이진트리 형태로 데이터를 정렬한다.해당 구조는 매 시간 바뀌는 전체 데이터 내에서 최댓값과 최솟값을 빠르게 찾기위해 사용된다.
해당 자료구조를 통해 탐색한 원소를 최솟값부터 가져와, 현재보다 작은 원소 값이 있으면 그 수를 저장하고,힙 구조에 계속해서 집어넣으면서 진행한다.
이때 자료를 정렬하는 heap 알고리즘의 시간복잡도는 NlogN이다.
따라서 해당 알고리즘을 활용해서 NlogN 시간내에 문제를 해결할 수 있다.
코드 내용은 아래와 같다.
import heapq
def solution(numbers):
answer = [-1] * len(numbers) # 기본적으로 모든 값은 -1로 초기화
max_heap = [] # 최대 힙을 사용하여 큰 값을 추적하기 위한 리스트
# 오른쪽에서부터 왼쪽으로 탐색
for i in range(len(numbers) - 1, -1, -1):
num = numbers[i]
# 현재 숫자보다 큰 숫자를 찾을 때까지 힙에서 제거
while max_heap and max_heap[0] <= num:
heapq.heappop(max_heap)
# 만약 힙에 남아있다면, 그것이 오른쪽에서의 첫 번째 큰 수
if max_heap:
answer[i] = max_heap[0]
# 힙에 현재 숫자를 넣음 (음수로 넣어 최대 힙처럼 동작하게 함)
heapq.heappush(max_heap, num)
return answer파이썬에서도 힙 알고리즘을 잘 활용할 수 있도록
익혀두는 것이 좋을 것 같다.
이어서, 그렇다면 다른 방법은 뭐가 있을까?
2. Stack 자료구조를 활용한 방법
다음 방법은 제목과 같이 Heap이 아닌, Stack 구조를 활용하는 방법이다.
모든 값을 순회하는 과정에서 그 값들을 모두 Stack 구조에 넣는다.
Stack이란, 후에 알고리즘으로 설명하겠지만
First In Last Out의 구조를 갖는 자료형이다.
말 그대로, 들어오는 데이터를 순서대로 쌓는 방식으로 저장하는 것이다. (Stack)
뺄 때는 위에서 부터 순차적으로 빼므로, 가장 먼저 들어간 데이터는 마지막에 나오게 된다.
그림을 통해서 Stack의 구조를 간단히 이해하고 넘어가자
결론적으로, 순회하며 모든 값을 Stack에 넣고, 현재 값과 비교하면서 Stack에서 하나씩 빼는 방식이다.
설명이 조금 부족한데,
앞선 heap으로 해결한 방법과 동일한 구조를 갖는다.
하지만 Stack에서 pop하는 과정은, 데이터를 빼내는 것을 의미하는데,
빼낸 데이터의 값을 return한다.
이를 활용하면 더욱 짧고 간결한 코드를 작성할 수 있다.
(참고로 파이썬에서 stack을 구현할 때는, List 자료형의 append함수와 pop함수를 사용하면 된다.)
def solution(numbers):
stack = []
result = [-1] * len(numbers)
for i in range(len(numbers)):
while stack and numbers[stack[-1]] < numbers[i]:
result[stack.pop()] = numbers[i]
stack.append(i)
return result
-1로 정답 리스트를 초기화 하고,
Stack에 차곡차곡 담으며, 해당 값이 탈출할 때,
그 값에 해당하는 큰 수 값을 정답 리스트에 입력한다.
해당 문제를 다양한 자료구조를 활용하여 해결할 수 있는 점이 매력적이었던 것 같다.