[고득점 Kit(스택/큐)] 주식가격 (Python)

 

코딩테스트 연습 - 주식가격

 

 

풀이

 입력으로 주어진 prices를 차례대로 스택에 push한다. 스택에는 (가격, 스택에 push된 시간)과 같은 튜플 형태로 넣어줬다. 여기에서 시간은 인덱스와 같은 의미를 지닌다. 스택에 넣기 전에 스택의 top에 있는 주식의 가격과 비교한다. 현재 스택의 top에 있는 주식 가격보다 현재 주식의 가격이 낮을 경우에는 스택에서 해당 주식을 pop을 통해 제거한다. 제거하면서 answer 배열에 '현재 시간 - 스택에 push된 시간'을 구하여 알맞은 인덱스에 넣어주면된다.

 

코드

# (가격, 들어온 시간)
 
def solution(prices):
    # 일단 answer 배열을 prices의 개수만큼 0으로 초기화해준다.
    answer = [0] * len(prices)
 
    st = []
    time = 0 # 현재 시간 
    idx = 0  # 해당 주식의 인덱스 (시간)
 
    for price in prices:
        # 현재 스택의 top에 있는 주식의 가격보다 현재 가격이 낮을 경우에 스택을 pop해준다.
        # time - popIdx를 통해 해당 주식의 가격이 몇 초만에 낮아졌는지를 구해서 answer 배열을 수정한다.
        while len(st) != 0 and st[-1][0] > price:
            (_, popIdx) = st.pop()
            answer[popIdx] = time - popIdx
 
        st.append((price, idx))
 
        idx += 1
        time += 1
 
    time -= 1
 
    # 현재 단계에 스택에 남아 있는 주식들은 끝날때까지 가격이 떨어지지 않은 주식이다. 
    # 남은 주식들의 가격이 떨어지지 않은 시간을 구하여 answer 배열을 수정한다.
    while st:
        (_, popIdx) = st.pop()
        answer[popIdx] = time - popIdx
 
    return answer