(프로그래머스 스택/큐 문제 풀이) 주식가격

스택/큐(Stack/Queue)

주식가격

문제 설명
초 단위로 기록된 주식가격이 담긴 배열 prices가 매개변수로 주어질 때, 가격이 떨어지지 않은 기간은 몇 초인지를 return 하도록 solution 함수를 완성하세요.

제한사항
prices의 각 가격은 1 이상 10,000 이하인 자연수입니다.
prices의 길이는 2 이상 100,000 이하입니다.

입출력 예
prices                      return
[1, 2, 3, 2, 3]         [4, 3, 1, 1, 0]

입출력 예 설명
1초 시점의 ₩1은 끝까지 가격이 떨어지지 않았습니다.
2초 시점의 ₩2은 끝까지 가격이 떨어지지 않았습니다.
3초 시점의 ₩3은 1초뒤에 가격이 떨어집니다. 따라서 1초간 가격이 떨어지지 않은 것으로 봅니다.
4초 시점의 ₩2은 1초간 가격이 떨어지지 않았습니다.
5초 시점의 ₩3은 0초간 가격이 떨어지지 않았습니다.

def solution(prices):
    answer = [0] * len(prices)
    price_stack = []
    
    for i, price in enumerate(prices):
        while price_stack and price < prices[price_stack[-1]]:
            j = price_stack.pop()
            answer[j] = i - j
        price_stack.append(i)
    
    while price_stack:
        j = price_stack.pop()
        answer[j] = len(prices) - 1 - j
    return answer

문제를 보고 딱 들었던 생각은 prices의 인덱스를 따로 뽑아서 가격이 몇 초 유지되었는지 판단해야겠다는 것이다. 그런데 해시 문제가 아닌 스택/큐 문제기 때문에 딕셔너리는 아닐 것이고, 그렇다면 어떻게 해야하는지 고민하다 결국 떠올리지 못했다.

해결 방법은 바로 prices와 길이가 똑같은 리스트를 만들고 인덱스에 맞춰서 가격이 몇 초 유지되었는지 정보를 기입하는 것이다. 이를 위해선 prices를 순환할 때 인덱스를 뽑아오도록 enumerate 함수를 사용해야 한다.

그리고 핵심은 새로운 price가 들어올 때 기존 price보다 크다면 stack에 추가해주고, 작다면 stack.pop()을 통해 가격이 유지된 시간을 기록한다. 그리고 끝까지 남아 있는 가격에 대해선 별도의 반복문 작업을 거쳐 시간을 기록해준다.

코드만 봐서는 이해가 안될 수 있으니 아래 과정을 살펴보자.

prices의 원소 2가 들어오는 순간 price stack에서 pop()이 호출되면서 answer[2]는 1초로 업데이트되었다.

while문을 통해 price stack의 남은 원소에 대해서도 유지된 시간을 구하면 최종적으로 아래와 같아진다.

다른 코드를 보면 brute force나 deque 방식으로도 풀었지만 시간 복잡도가 $O(N^2)$가 되며, 위의 코드로 하면 $O(N)$이 보장된다는 점에서 best가 아닐까 싶다.