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

문제 설명

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

제한사항

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

입출력 예pricesreturn

[1, 2, 3, 2, 3]

[4, 3, 1, 1, 0]

풀이

스택을 활용하여 풀었다. 스택에 집어넣으면서, 스택 맨위값과 비교하면서 감소하면(stack[num_element-1][1] > price 상황)스택에서 pop하며 현재의 시간에서 price의 시간을 뺴서(idx - ith) 떨어지지 않은 기간을 기록한다.

for문을 다 돌고, stack에 남아있다면(주식 가격이 떨어지지 않았다면) 마지막 시간(len(price)-1)과 주식의 시간(idx)을 빼서 기록해준다.

def solution(prices):
    answer =[0] * len(prices)
    stack = []
    num_element = 0 
    for idx,price in enumerate(prices) :
        while stack and stack[num_element-1][1] > price :
            # print(stack[num_element-1][1],"vs", price)
            ith = stack[num_element-1][0]
            answer[ith] = idx - ith
            stack.pop()
            num_element -=1
            # print(stack,answer)
        stack.append([idx,price])
        num_element +=1
        # print(idx,stack,answer)
    while stack :
        idx,price = stack.pop()
        answer[idx] = len(prices)-1 - idx
    # print(stack,answer)
    return answer