[프로그래머스/Python] 뒤에 있는 큰 수 찾기

def solution(numbers):
    answer = [-1]
    stack = [numbers[-1]]
    
    for i in range(len(numbers)-2, -1, -1):
        while stack and stack[-1] <= numbers[i]: # 같아도 pop
            stack.pop()
        
        if stack:
            answer.append(stack[-1])
        else:
            answer.append(-1)
            stack.append(numbers[i])
        
        if numbers[i] > numbers[i-1]:
            stack.append(numbers[i])
    
    answer.reverse()
    return answer

백준-탑 문제와 유사하다.

stack에 큰 숫자들을 쌓는다. 현재 top보다 큰 숫자가 나오면 pop()을 해준다.

스택이 비어있으면, 현재 숫자보다 큰 숫자가 뒤에 없었다는 뜻이므로 -1을 넣어준다.

스택이 남아 있다면, 현재 숫자의 뒤에 있는 큰 수는 top 이다.

 

그리고 마지막에 현재 값도 넣어주는데, 현재 값이 다음 값보다는 큰 수일 수 있기 때문이다.

(사실 if문 없어도 될 것 같다. 작으면 알아서 pop 될거니까)

 

for문을 반대로 돌아서 답을 넣어줬으니, 답도 반대로 출력해주면 된다.

 

def solution(numbers):
    n = len(numbers)
    stack = [] # idx 스택
    answer = [-1] * n
    
    for i in range(n):
        while stack and numbers[i] > numbers[stack[-1]]:
            idx = stack.pop()
            answer[idx] = numbers[i]
        stack.append(i)
    return answer

탑에서 익힌 풀이를 그대로 써도 정답이다!

idx를 이용한 스택을 사용한다.

스택의 top보다 큰 숫자를 만났을 때, 그 동안의 index의 answer에 해당 숫자를 넣어주는 방법! (몰아서 처리랄까?)

 

굳이 한번 더 코드를 줄이자면 아래와 같다

answer[stack.pop()] = numbers[i]