[leetcode #309] Best Time to Buy and Sell Stock with Cooldown

Problem

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

Idea

오랜만에 재미있는 문제가 나왔다.

주어진 주식 한 종목이 매일 가격이 변동하는데, buyer는 사고 팔거나 거래를 하지 않을 수 있다. 대신 주식을 팔았다면 다음 날에는 거래를 할 수 없는 제한 사항이 있다.

Buyer는 주식을 보유하고 있거나 상태거나, 주식을 가지고 있지 않은 상태 둘 중 하나다. 주식을 매매 또는 매도할 때 가지고 있는 현금을 각 상태에 저장하고 마지막에 주식을 매도한 상태에서 보유할 수 있는 현금의 최대값을 답으로 주면 된다.

두 상태를 저장하는 dp array를 buy와 cooldown으로 정했다. buy array는 이전 날짜에서 주식을 보유하고 있지만 매도하지 않은 경우, 주식을 보유하고 있지 않다가 당일에 매매한 경우 두 가지가 있다. 둘 중 가지고 있는 현금이 더 큰 값을 dp에 저장하면 된다. 당연하게도 이는 주식을 더 싼 값에 사는 것이 이득이기 때문이다.

cooldown array는 보유한 주식을 매도한 경우와 주식을 보유하지 않은 상태로 매매하지 않는 경우 두 가지고 있다. 대신 위의 제약조건에 의해 주식을 매도하면 다음 날에 매매를 하지 못하므로 다음 날에 주식을 보유하지 않은 상태에 가지고 있는 현금을 저장한다. 주식을 매매하지 않는 경우는 이전의 현금과 현재의 현금 (이틀 전 주식을 매도한 경우)를 비교해 더 큰 값을 저장하면 된다.

prices array 전체를 탐색한 뒤 당일에 주식을 보유하지 않은 상태와 다음 날 주식을 보유하지 않은 상태를 비교해 더 큰 값을 리턴하면 된다. 주식 거래가 가능한 날 다음날까지 고려하는 이유는 마지막 날에 주식을 판 경우까지 고려해야 하기 때문이다.

Solution

class Solution {
    public int maxProfit(int[] prices) {
        int n = prices.length;
        int[] buy = new int[n+1];
        int[] cooldown = new int[n+1];

        buy[0] = -prices[0];
        cooldown[0] = 0;

        for (int i=1; i < n; i++) {
            buy[i] = Math.max(buy[i-1], cooldown[i-1] - prices[i]);
            cooldown[i+1] = buy[i-1] + prices[i];
            cooldown[i] = Math.max(cooldown[i-1], cooldown[i]);
    }
   
        return Math.max(cooldown[n-1], cooldown[n]);
    }
}

어떻게 하면 더 적은 메모리와 0ms로 풀 수 있는지 궁금하다.

Reference

https://leetcode.com/problems/best-time-to-buy-and-sell-stock-with-cooldown/