[프로그래머스 코딩테스트 연습] 주식가격 (Python3)

https://school.programmers.co.kr/learn/courses/30/lessons/42584

def solution(prices):
    N = len(prices)
    answer = [0 for _ in range(N)]
    stack = []
    
    for now, price in enumerate(prices):
        # 스택 인덱스 가격과 현재 가격 비교. 
        # 마지막 가격이 현재가격보다 크면 가격 떨어진 것이므로 인덱스 꺼냄
        while stack and prices[stack[-1]] > price:
            idx = stack.pop()
            answer[idx] = now - idx # 현재 인덱스 - 꺼낸 인덱스 
        stack.append(now) # 현재 인덱스 추가
        
    while stack:
        idx = stack.pop()
        answer[idx] = N-idx-1 # 끝까지 남은 경우
    return answer

 

요약

1. 스택에는 인덱스만 저장한다.

2. 가격이 떨어질 때 마다 스택에서 값을 꺼내 시간을 계산해 answer에 저장

3. 최종적으로 스택에 남은 값은 끝까지 유지된 시간을 계산한다.

---

prices = [1,2,3,2,3] 일 때의 단계를 살펴보자

 

초기 상태 : stack = [], answer = [0,0,0,0,0]

 

1초 (now=0, price=1) : stack = [0], answer = [0,0,0,0,0]

2초 (now=1, price=2) : stack = [0, 1], answer = [0,0,0,0,0]

3초 (now=2, price=3) : stack = [0, 1, 2], answer = [0,0,0,0,0]

 

4초 (now=3, price=2) : 가격이 떨어지는 구간

stack[-1] = 2으로, 스택 마지막 가격 prices[2] = 3 > 현재 price = 2 조건 진입. 

idx = 2. answer[2] = 3 - 2 = 1

stack = [0, 1, 3], answer = [0,0,1,0,0]

 

5초 (now=4, price=3) : stack = [0, 1, 3, 4], answer = [0,0,1,0,0]

 

while stack: 구간을 통해 종료 후 스택을 처리한다.

(최종적으로 스택에 남은 값은 끝까지 유지된 시간을 계산하게 된다.)

idx=4, answer[4] = 5-4-1 = 0, answer = [0,0,1,0,0]

idx=3, answer[3] = 5-3-1 = 1, answer = [0,0,1,1,0]

idx=1, answer[1] = 5-1-1 = 3, answer = [0,3,1,1,0]

idx=0, answer[0] = 5-0-1 = 4, answer = [4,3,1,1,0]