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

作者 Shilei Tian 日期 2017-05-02
[LeetCode 309] Best Time to Buy and Sell Stock with Cooldown

题目描述

Say you have an array for which the ith element is the price of a given stock on day i.

Design an algorithm to find the maximum profit. You may complete as many transactions as you like (ie, buy one and sell one share of the stock multiple times) with the following restrictions:

  • You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).
  • After you sell your stock, you cannot buy stock on next day. (ie, cooldown 1 day)

Example:

prices = [1, 2, 3, 0, 2]
maxProfit = 3
transactions = [buy, sell, cooldown, buy, sell]

解题思路

这是一道很典型的动态规划的题目,但是状态转移感觉不是一眼就能看出来。LeetCode 上面这道题的 Solutions 中排名第二的是 Share my DP solution (By State Machine Thinking),它利用了状态机的思想,表明了不同状态的转移情况,这是我第一次在解决 DP 问题的时候看到有这么用,值得记录下来。

根据题意,我们可以推断出,一共有三种状态:未买入状态(S0)、已买入状态(S1)和已卖出状态(S2),所谓的已卖出,即题目要求的冷却。我们可以从 S0 通过买入进入 S1,也可以 S0 保持不动(即不做任何买卖操作)还是进入 S0;同理,S1 可以不进行任何操作,状态还是转移到 S1,以及通过卖出操作转移到 S2。对于 S2 而言,它只有一个状态转移,即冷却完毕后进入 S0,状态转移图如下所示:

State Transition

有了状态转移,我们就可以定义 DP 的结构了。这道题目一个精彩的地方,就在于定义三个数组,分别记录这三种状态下在任何时间可以达到最大的收益,我们记这三个数组分别为 s0s1s2。那么如何更新这三个数组呢?就根据状态转移图来:

s0[i] = max(s0[i - 1], s2[i - 1])
s1[i] = max(s1[i - 1], s0[i - 1] - prices[i])
s2[i] = s1[i - 1] + prices[i]

对于每一个状态而言,有多少条入边,max 函数就有多少个参数。我们拿 s1 来说,它在时间点 i 的收益就等于下面两个值的最大值:

  1. s1[i - 1],对应的操作为保持不变,不采取任何措施;
  2. s0[i - 1] - prices[i],对应的操作为买入,这时需要用上一时间点在未买入状态的收益减当前时间点的价格。

那么初始值应该如何设定呢?对于 s0 而言,最开始收益为 0,这个比较显然;对于 s1 而言,一开始我们就执行了购买操作,因此它最开始的收益就等于 -prices[0],而对于 s[2] 来说,我们没有执行卖出操作,所以它刚开始就等于 INT_MIN 就好了。

这样就有了下面的程序:

class Solution {
public:
int maxProfit(vector<int>& prices) {
if (prices.size() < 2) {
return 0;
}
vector<int> s0(prices.size(), 0);
vector<int> s1(prices.size(), 0);
vector<int> s2(prices.size(), 0);
s0[0] = 0;
s1[0] = -prices[0];
s2[0] = INT_MIN;
for (int i = 1; i < prices.size(); ++i) {
s0[i] = max(s0[i - 1], s2[i - 1]);
s1[i] = max(s1[i - 1], s0[i - 1] - prices[i]);
s2[i] = s1[i - 1] + prices[i];
}
return max(s0[prices.size() - 1], s2[prices.size() - 1]);
}
};