[Programmers] 주식가격 (Python)

주식가격

programmers.co.kr/learn/courses/30/lessons/42584

 

알고리즘 문제 풀이에 자주 보이는, next greater element 랑 비슷한 문제다.

주어진 배열에서 각 인덱스가 i < j 일때 A[i] 와 A[j] 간을 비교하여 j - i 만큼을 정답에 추가하는 문제다.

 

빠르게 하려고 파이썬으로 풀어보았다. (파이썬에서 list 는 스택과 동일하므로 list 를 사용하였다.)

def solution(prices):
    # 정답의 전체 길이는 prices 와 동일하다. 그래서 prices 길이만큼 0 으로 초기화 시켜주었다.
    answer = [0 for _ in range(len(prices))]
    stack = list()
    # 0 번쨰 값은 미리 스택에 넣어두었다.
    stack.append(0)
    
    # prices 를 인덱스 접근으로 순회
    for index in range(1, len(prices)):
        # 스택에 값이 있고, 스택의 가장 마지막 (last) 값이 현재 price 보다 크다면 주식 값이 떨어진 것이다.
        # 스택의 가장 마지막 값은 인덱스 값이므로 prices[stack[-1]] 로 비교하였다.
        while stack and prices[stack[-1]] > prices[index]:
            # 현재 price 보다 큰 주식 을 찾는다.
            # 바꿔 말하면 index 시점에서 주식이 떨어졌다고 봐도 되는 index 이전의 모든 주식을 찾는다는 얘기다.
            # index 이전의 모든 주식은 stack 에 넣어놨기 때문에 stack 을 반복문으로 돈다.
            prev = stack.pop()
            # prev 는 index 이전의 어느 값이다
            # 조건에선 단순 비교를 했고 여기서 실제로 pop 을 해 준다. (파이썬에선 자바의 poll, peek 메서드가 없기 때문)
            # prev 시점에 샀던 주식 > index 시점의 주식이 된다.
            # 이 인덱스 간의 차이가 기간 (초) 가 된다.
            answer[prev] = (index - prev)
        
        # 현재 주식에 해당하는 index를 넣어준다.
        stack.append(index)
    
    # 여기까지 왔으면 prices 를 전부 순회하였다.
    for i in range(len(answer)):
        # answer[i] 가 0 이라는 의미는 i 시점에 한번도 주식이 떨어지지 않았다는 것이다.
        if answer[i] == 0:
            # 그러면 len(prices) 만큼의 시간 동안 끝까지 떨어지지 않았다는 의미이므로
            # 전체 prices 길이 - i - 1 을 업데이트 해준다.
            answer[i] = len(prices) - i - 1
    return answer

 

최대한 주석에 풀이를 작성하려고 하였다.

 

정답 answer 는 전체 배열의 크기와 동일하므로 0으로 prices 사이즈만큼 초기화 시켜주었다. 그리고 첫번째 원소 값은 미리 스택에 넣어두었다.

 

스택에 있는 값은 인덱스과 각 인덱스에 해당하는 prices 값으로, prices 를 순회하면서 스택에 있는 last 값과 비교를 하게 된다.

 

스택에 있는 가장 마지막 값은 index 시점 보다 이전의 값으로 index - x (코드에선 prev) 에 해당하는 어떤 지점일 것이다.

이 값과 현재 prices 를 비교하여 주식이 떨어졌는지 보고, 떨어졌다면 answer 에 각 인덱스 차이 만큼을 넣어준다.

 

스택을 계속 순회하는 이유는 [1, 3, 3, 2, 2] 같은 예제로 보자면, 2번째 주식과 3번째 주식을 비교하여 2번째 주식이 떨어졌음은 확인 가능하지만 1번째 주식이 3번째 시점에 떨어졌는지는 모르기 때문이다. 스택에 있는 값의 가장 마지막 값을 계속 pop 하여서 비교해야 한다.

 

예제에서 [1, 2, 3, 2, 3] 가 주어졌을 때, prices 를 다 순회한 시점에서 answer 는 [0, 0, 1, 0, 0] 이 된다.

 

그리고 나머지 0 으로 되어 있는 부분에 해당하는 인덱스의 prices 는 한 번도 주식 이 떨어지지 않았다는 것이다. 그래서 전체 prices 길이 5에서 각 인덱스 -1 을 빼 주면 그 값에 해당하는 초 만큼 떨어지지 않았다고 볼 수 있다.

 

말이 좀 어렵지만, 결국 answer[0] 은 전체  prices 길이 5 - 0 - 1 을 했고 그 값인 4가 나왔다는 것이다. 그러면 answer[0] 은 4초동안 주식이 떨어지지 않았다는 것을 의미한다.