시간 순서대로 정렬된 주식 가격 리스트 prices가 주어졌을 때, 주식을 최대 두 번 사고팔아 얻을 수 있는 최대 수익을 구하는 문제를 살펴보겠습니다. 단, 반드시 먼저 매수한 후에 매도해야 하며, 두 번째 매수는 첫 번째 매도 이후에 이루어져야 합니다.
문제 예시
입력이 prices = [2, 6, 3, 4, 2, 9]라면 출력은 11이 됩니다. 가격 2에 매수한 후 6에 매도하고(+4), 다시 가격 2에 매수한 뒤 9에 매도하면(+7) 총 수익이 11이 되기 때문입니다.
풀이 전략
이 문제는 네 개의 상태 변수를 이용해 한 번의 순회로 해결할 수 있습니다. 각 변수는 다음 상태를 추적합니다.
- first_buy: 첫 번째 매수 후 보유 중인 상태의 최대 이익
- first_sell: 첫 번째 매도까지 완료한 상태의 최대 이익
- second_buy: 두 번째 매수 후 보유 중인 상태의 최대 이익
- second_sell: 두 번째 매도까지 완료한 상태의 최대 이익
각 가격을 순회하면서 다음과 같이 값을 갱신합니다.
- first_buy := max(first_buy, -px)
- first_sell := max(first_sell, first_buy + px)
- second_buy := max(second_buy, first_sell - px)
- second_sell := max(second_sell, second_buy + px)
모든 가격을 확인한 후에는 0, first_sell, second_sell 중 최댓값을 반환합니다. 음수가 나오는 경우(즉, 어떤 거래도 이익이 되지 않는 경우)를 대비해 0을 포함시키는 것입니다.
구현 코드
class Solution:
def solve(self, prices):
first_buy = first_sell = float("-inf")
second_buy = second_sell = float("-inf")
for px in prices:
first_buy = max(first_buy, -px)
first_sell = max(first_sell, first_buy + px)
second_buy = max(second_buy, first_sell - px)
second_sell = max(second_sell, second_buy + px)
return max(0, first_sell, second_sell)
ob = Solution()
prices = [2, 6, 3, 4, 2, 9]
print(ob.solve(prices))
입력
[2, 6, 3, 4, 2, 9]
출력
11
정리
이 알고리즘은 리스트를 한 번만 순회하므로 시간 복잡도는 O(n), 추가 메모리는 상수 개수의 변수만 사용하므로 공간 복잡도는 O(1)입니다. 매수·매도 횟수 제한이 k번으로 일반화된 경우에도 같은 아이디어를 확장하여 적용할 수 있습니다.