[Python] 프로그래머스 주식가격

https://school.programmers.co.kr/learn/courses/30/lessons/42584?language=python3#

프로그래머스

풀이

초 단위로 기록된 주식 가격이 주어질 때, 가격이 떨어지지 않은 기간이 몇초인지 리턴해야 하는 문제입니다.

 

스택에 값을 기록한뒤 새로운 값이 들어오면, 이전에 기록한 값이 새로운 값 보다 크다면 계속해서 pop해주는 방법으로 문제를 해결할 수 있습니다.

이 때, pop 하는 값은 몇초간 스택에 남아있었는지 answer에 기록해둡니다.

 

값을 검사하고 stk에 남아있는 값들에 대해 해당 인덱스가 stk에 몇초간 남아있었는지 채워줍니다.

이 때, stk에 머무른 시간은 len(prices) - 1 - 주식의 인덱스 입니다.

 

참고

 pop하는 부분을 while문으로 처리하지 않으면 이전에 넣은 값에 대해서는 처리할 수 없음에 주의가 필요합니다.

예를 들어,

[3, 3, 2, 2, 2] 의 경우 [2, 1, 2, 1, 0]이 정답이나, 한번만 처리하는 로직으로 짠다면 [4, 1, 2, 1, 0] 이라는 값이 나오게 됩니다. 

def solution(prices):
    answer = [0 for _ in prices]
    stk = []
    for idx, v in enumerate(prices):
        if idx == 0:
            stk.append([idx, v])
        else:
            # 가격이 떨어진 경우, 지금까지 넣어둔 모든 값을 비교해야함
            while stk and stk[-1][1] > v:
                answer[stk[-1][0]] = idx - stk[-1][0]
                stk.pop()
            
            stk.append([idx, v])
            
    # 아직 인덱스가 없는 시점 채우기
    for v in stk:
        answer[v[0]] = len(prices) - 1 - v[0]
        
    return answer