프로그래머스 > 뒤에 있는 큰 수 찾기 (Python 3)

1. 문제 설명

정수로 이루어진 배열 numbers가 있다.

배열에서 자신보다 뒤에 있는 숫자 중에서, 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 한다.

정수 배열 numbers가 매개변수로 주어질 때, 모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 반환하라.

단, 뒷 큰수가 존재하지 않는 원소는 -1을 담는다.

 

예시)

- numbers = [2, 3, 3, 5]

정답 : [3, 5, 5, -1]

 

- numbers = [9, 1, 5, 3, 6, 2]

정답 : [-1, 5, 6, 6, -1, -1]

 

 

 

 

2. 코드

(1) 1차 시도

def solution(numbers):
    answer = []
    
    for i in range(len(numbers)):
        for j in range(i+1, len(numbers)):
            if numbers[i] < numbers[j]:
                answer.append(numbers[j])
                break
        if len(answer) != (i+1):
            answer.append(-1)
    
    return answer

 

 

 

(2) 2차 시도

def solution(numbers):
    answer = [-1 for _ in range(len(numbers))]
    stack = []
    
    for i in range(len(numbers)):
    
        while stack and numbers[stack[-1]] < numbers[i]:
            answer[stack.pop()] = numbers[i]
        
        stack.append(i)
        
    return answer

 

 

 

3. 설명

이 문제가 되게 단순해 보였지만, 시간 초과를 극복하려면 새로운 관점이 필요하기 때문에 어려웠다.

 

많은 사람들의 코드를 봤을 때, 2차 시도 코드와 비슷하거나 동일하다.

그래서 나는 어떻게 stack을 도입할 시도를 하게 되었는지를 생각해보기로 했다.

처음에는 절대로 저런 생각이 나지 않기 때문이다..

 

 

index 하나 하나씩을 해결해나가는 식이 아니라, 바로 뒤에 큰 수가 없다면 잠시 스택에 넣고 대기하는 형식이다.

시간 초과가 나는 방식

이렇게 하나씩 뿌시는 경우에는 현재 인덱스에서 그 뒤의 인덱스 이후를 다 비교해봐야 하기 때문에, $O(n^2)$이 된다.

 

 

검토하고 있는 인덱스의 그 다음 인덱스 시점에서 바라본다.

만약 그 인덱스가 큰 수를 찾지 못했다면 스택에 넣는다!

 

기본적으로 현재 인덱스를 스택에 넣어주고, 조건 맞을 시에 pop() 하도록 구현한다.

위의 그림이 단계를 거치면서 스택의 변화를 보여준다.

 

즉, for 문이 돌아가는 그 시점의 인덱스는 answer 값이 변하는 시점의 인덱스가 아니라, 비교대상이라고 생각하면 된다!