파이썬으로 주식 수익을 내기 위해 기다려야 하는 최소 일수 구하기

시간 순서대로 나열된 어떤 회사의 일별 주가를 담고 있는 가격 리스트가 있다고 가정해 보겠습니다. 이때 원본 리스트와 길이가 같은 새로운 리스트를 만들어야 하며, 새 리스트의 인덱스 i에 해당하는 값은 해당 날짜부터 이익을 얻을 때까지 기다려야 하는 최소 일수입니다. 만약 어떤 방법으로도 이익을 얻을 수 없다면 그 값은 0이 됩니다.

예를 들어 입력이 prices = [4, 3, 5, 9, 7, 6]이라면 출력은 [2, 1, 1, 0, 0, 0]이 됩니다. 첫째 날의 가격 4는 이틀 뒤 가격 5에서 수익을 낼 수 있으므로 2가 되고, 마지막 세 날은 이후 더 높은 가격이 없으므로 0이 됩니다.

문제 해결 접근 방식

이 문제는 '단조 감소 스택(Monotonic Decreasing Stack)'을 활용하면 효율적으로 해결할 수 있습니다. 단계별로 살펴보면 다음과 같습니다:

• ans := prices와 같은 크기의 리스트를 만들고 모든 값을 0으로 초기화합니다.
• q := 새로운 빈 리스트(스택)를 생성합니다.
• prices의 각 인덱스 i와 가격 p에 대해 다음을 반복합니다.
  • q가 비어 있지 않고, p가 q의 마지막 항목에 저장된 가격보다 클 동안:
    • j := q의 마지막 요소에 저장된 인덱스
    • ans[j] := i - j (현재 인덱스에서 과거 인덱스를 뺀 대기 일수)
    • q에서 마지막 요소를 삭제합니다.
  • (i, p) 쌍을 q의 끝에 삽입합니다.
• 모든 반복이 끝나면 ans를 반환합니다.

이 알고리즘의 시간 복잡도는 O(n)입니다. 각 인덱스는 스택에 최대 한 번 삽입되고 한 번 제거되기 때문에 전체 연산 횟수가 입력 크기에 비례합니다.

아래 예제 코드를 통해 더 자세히 이해해 보겠습니다:

예제 코드

class Solution:
   def solve(self, prices):
      ans = [0 for _ in prices]
      q = []
      for i, p in enumerate(prices):
         while q and p > q[-1][1]:
            j = q[-1][0]
            ans[j] = i - j
            q.pop()
         q.append((i, p))
      return ans

ob = Solution()
prices = [4, 3, 5, 9, 7, 6]
print(ob.solve(prices))

입력

[4, 3, 5, 9, 7, 6]

출력

[2, 1, 1, 0, 0, 0]