[algorithm] Best Time to Buy and Sell Stock with Cooldown

오현우·2022년 7월 2일
0

알고리즘

목록 보기
22/25
post-thumbnail

문제 내용

You are given an array prices where prices[i] is the price of a given stock on the ith day.

Find the maximum profit you can achieve. You may complete as many transactions as you like (i.e., buy one and sell one share of the stock multiple times) with the following restrictions:

After you sell your stock, you cannot buy stock on the next day (i.e., cooldown one day).
Note: You may not engage in multiple transactions simultaneously (i.e., you must sell the stock before you buy again).

Example 1:

Input: prices = [1,2,3,0,2]
Output: 3
Explanation: transactions = [buy, sell, cooldown, buy, sell]
Example 2:

Input: prices = [1]
Output: 0

Constraints:

1 <= prices.length <= 5000
0 <= prices[i] <= 1000

처음 접근 방법

profit = buy - curr_price 라고할 때
[1, 3, 10, 100, 0, 91, 1000] 의 예시로 보면 결국엔 변곡점은 100 > 0에서 갈 때 profit이 음수가 된다. 따라서 이때 우리가 주식을 100에 팔것인지, 10에 팔아서 91에 사서 이득을 볼 것인지 판단하여야 한다.

결국 여기서 내가 가장 중요하게 생각한 것은 음수로 profit이 바뀌기 이전 즉 10 > 100 이득이 0 91보다 커야만 100에 파는 것이 이득이라고 생각할 수 있게 된다.

따라서 해당 엣지 케이스만 고려해서 코드를 아래와 같이 작성했다.

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        curr_profit = prev_profit = answer = 0
        buy = prices[0]
        sell = False
        
        for price in prices[1:]:
            curr_profit = price - buy
            if curr_profit >= 0 and not sell:
                answer += curr_profit
                prev_profit = curr_profit
                buy = price
                
            elif curr_profit >= 0 and sell:
                if curr_profit > prev_profit:
                    answer += curr_profit - prev_profit
                sell = False
                buy = price
                prev_profit = curr_profit
                
            else:
                 buy = price
                 sell = True

        
        return answer

하지만 문제가 발생했다.
[1, 3, 10, 100, 91, 0, 1000]를 고려해보자.

우리는 모노토닉 즉 증가하는 수열에 대해서만 고려하였었고 지그재그로 발생하는 수열에 대해서는 대처하지 못했다.

나는 이러한 문제의 원인을 케이스를 정확히 구분하지 않고 조건문을 남발해서 발생하는 문제라고 생각하였다.

이러할 때에는 다른 사람의 코드를 보면서 리프레쉬 하는 것이 좋다.

dp를 활용한 풀이

결국 우리는 3가지의 상태의 중복없는 순열의 case로 나눌 수 있다.

3가지 상태: hold | not_hold | cooldown

3 * 2 = 6

6개의 케이스 이지만 cool_down > hold 즉 휴식기에 매수할 수 없으므로 5가지의 케이스로 나눌 수 있다.

hold >>>> do nothing >>>> hold
hold >>>> sell >>>> cooldown
not_hold >>>> do nothing >>>> not_hold
not_hold >>>> buy >>>> hold
cool_down >>>> do nothing >>>> not_hold

위의 케이스들을 기반으로 kadene's algo를 사용하여 답을 도출하면 아래와 같다.

class Solution:
    def maxProfit(self, prices):
        notHold, notHold_cooldown, hold = 0, float('-inf'), float('-inf')
        for p in prices:
            hold, notHold, notHold_cooldown = max(hold, notHold - p), max(notHold, notHold_cooldown), hold + p
        return max(notHold, notHold_cooldown)
profile
핵심은 같게, 생각은 다르게

0개의 댓글