Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 주식 매수·매도를 통해 얻을 수 있는 최대 이익 계산하기

문제 개요

시간 순서대로 기록된 회사의 주가를 담고 있는 숫자 리스트 nums가 주어진다고 가정해 봅시다. 하루에 최대 한 주씩 매수할 수 있으며, 여러 주식을 동시에 보유하는 것도 가능하고, 원하는 날짜에 얼마든지 매도할 수 있습니다. 이때 얻을 수 있는 최대 이익을 반환하는 프로그램을 작성해야 합니다.

예시

입력이 nums = [3, 4, 7, 3, 5]라면 출력은 9가 됩니다.

이유는 다음과 같습니다. 가격이 3일 때와 4일 때 각각 한 주씩 매수한 뒤, 가격이 7일 때 두 주 모두 매도합니다. 이후 다시 가격이 3일 때 매수하고, 가격이 5일 때 매도하면 됩니다.

총 이익 = (7 − 3) + (7 − 4) + (5 − 3) = 9

풀이 접근 방법

이 문제는 뒤에서부터 주가를 확인하며 해결할 수 있습니다. 핵심 아이디어는 현재보다 낮은 가격에서 매수했다면 모두 현재 가격에 팔아 이익을 취하는 것입니다.

  • 총 이익을 저장할 변수 ans를 0으로 초기화합니다.
  • 리스트가 빌 때까지 다음 과정을 반복합니다.
    • 리스트의 마지막 요소를 꺼내 top에 저장합니다(현재 매도 가격).
    • 리스트가 비어 있지 않고, top이 남은 마지막 요소보다 큰 동안 다음을 수행합니다.
      • ans(top - 해당 요소)를 더합니다.
      • 해당 요소를 리스트에서 제거합니다.
  • 모든 반복이 끝나면 ans를 반환합니다.

Python 구현 예제

아래 코드를 통해 실제 동작을 더 잘 이해할 수 있습니다.

def solve(nums):
    ans = 0
    while nums:
        top = nums.pop()
        while nums and top > nums[-1]:
            ans += top - nums.pop()

    return ans

nums = [3, 4, 7, 3, 5]
print(solve(nums))

입력

[3, 4, 7, 3, 5]

출력

9

동작 원리 정리

위 알고리즘은 스택과 유사한 방식으로 리스트의 뒤쪽부터 탐색하면서, 현재 가격(top)보다 낮은 가격들을 하나씩 제거하며 그 차액만큼 이익에 누적합니다. 결국 각 매도 시점마다 가능한 모든 저가 매수분을 정산하는 것과 같으므로, 전체 탐색이 끝나면 자연스럽게 최대 이익이 완성됩니다. 시간 복잡도는 각 요소가 한 번씩만 추가되고 제거되므로 O(n)으로 효율적입니다.