description

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

submission

// we'll solve this by dynamic programming and state machine thinking
// we keep track of max profit of every moments and states
// we have 3 possible states
//  - invested: we have bought a stock
//  - cooldown: we just sold a stock
//  - uninvested: we can buy a stock immediately
impl Solution {
    pub fn max_profit(prices: Vec<i32>) -> i32 {
        // make sure we have at least one price
        if prices.is_empty() {
            return 0;
        }
        // init states with the 1st price
        let (mut invested, mut cooldown, mut uninvested) = (-prices[0], 0, 0);
        // skip the 1st price
        for &price in &prices[1..] {
            // we have the following state transition
            // uninvested --buy--> invested
            invested = invested.max(uninvested - price);
            // cooldown --wait--> uninvested
            uninvested = uninvested.max(cooldown);
            // invested --sell--> cooldown
            cooldown = cooldown.max(invested + price);
        }
        // max profit can't happen when we are in invested state 
        cooldown.max(uninvested)
    }
}