股票买卖状态机解决方案分析
题目:
给定一个整数数组 prices ,它的第 i 个元素 prices[i] 是一支给定的股票在第 i 天的价格。
设计一个算法来计算你所能获取的最大利润。你最多可以完成 k 笔交易。
穷举框架
看看总共有几种「状态」,再找出每个「状态」对应的「选择」。我们要穷举所有「状态」,穷举的目的是根据对应的「选择」更新状态。如下所示:
for 状态1 in 状态1的所有取值:
for 状态2 in 状态2的所有取值:
for ...
dp[状态1][状态2][...] = 择优(选择1, 选择2, ...)
具体到当前问题,每天都有3中「选择」,买入,卖出,无操作, 用buy, sell, rest 表示
这个问题的「状态」有三个,第一个是天数,第二个是当天允许交易的最大次数,第三个是当前的持有状态(即之前说的 rest 的状态,我们不妨用 1 表示持有,0 表示没有持有)
用一个三维数组 dp 就可以装下这几种状态的全部组合,用 for 循环就能完成穷举:
// dp[i][k][0|1] 第i天,最多进行k次交易,0|1 表示当前是否持有股票
// 0 <= i < n, 1 <= k <= K; n为天数,大K为最大交易数,此问题供n*K*2中状态,全部穷举就能搞定
for 0 <= i < n:
for 1 <= k <= K:
for s in {0, 1}:
dp[i][k][s] = max(buy, sell, rest)
状态转移框架
思考状态可以怎么转移:
0 -> rest -> 0; 0 -> buy -> 1; 1 -> rest -> 1; 1 -> sell -> 0;
根据上述转移关系,可以写出状态转移方程:
dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
= max( 选择 rest, 选择 sell)
解释:今天没有持有股票,有两种可能:
(1)昨天没有持有,今天选择rest, 依然没有持有
(2)昨天持有,今天选择sell,今天没有持有了
dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])
= max( 选择rest, 选择buy)
解释: 今天持有股票,有两种可能:
(1)昨天持有,今天选择rest,依然持有
(2)昨天没有,今天选择buy持有了
注意: 如果buy,就要从利润中减去prices[i], 反之如果sell,就加上prices[i];
另外buy的时候交易次数要减去1,也可以在sell的时候减去1,效果是一样的
base case:
dp[-1][k][0] = 0; i 是从0开始的,-1表示还没有开始,利润为0
dp[-1][k][1] = -infinity; 还没有开始不可能持有股票,用-infinity表示这种不可能
dp[i][0][0] = 0; k 从1开始,k == 0 表示不允许交易,利润为0
dp[i][0][1] = -infinity; 不允许交易的情况下,不能持有股票,用-infinity 表示这种不可能
套用框架
-
k = 1
dp[i][1][0] = max(dp[i-1][1][0], dp[i-1][1][1] + prices[i]) dp[i][1][1] = max(dp[i-1][1][1], dp[i-1][0][0] - prices[i]) = max(dp[i-1][1][1], -prices[i]) 现在发现k都是1,不会改变,即k对状态转移没有影响了,可以进一步化简去掉所有k: dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i]) dp[i][1] = max(dp[i-1][1], -prices[i])
翻译成代码:
int n = prices.length; int[][] dp = new int[n][2]; for (int i = 0; i < n; i++) { dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]) dp[i][1] = Math.max(dp[i-1][1], -prices[i]) } return dp[n-1][0];
处理base case,且新状态只和前一个状态有关,空间优化:
int maxProfit_k_1(int[] prices) { int n = prices.length; // base case, dp[-1][0] = 0, dp[-1][1] = -infinity int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE; for (int i = 0; i < n; i++){ dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]); dp_i_1 = Math.max(dp_i_1, -prices[i]); } return dp_i_0; }
-
k = +infinity
如果k为无穷,那么可能认为k和k-1是一样的,可以改写框架:
dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i]) dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i]) = max(dp[i-1][k][1], dp[i-1][k][0] - prices[i]) 同理,k不会改变,不需要记录k这个状态: dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i]) dp[i][1] = max(dp[i-1][1], dp[i-1][0] - prices[i])
翻译成代码:
int maxProfit_k_inf(int[] prices){ int n = prices.length; int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE; for (int i = 0; i < n; i++){ int temp = dp_i_0; dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]); dp_i_1 = Math.max(dp_i_1, temp - prices[i]); } return dp_i_0; }
-
k = +infinity with cooldown
每次 sell 之后要等一天才能继续交易。只要把这个特点融入上一题的状态转移方程即可:
dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i]) dp[i][1] = max(dp[i-1][1], dp[i-2][0] - prices[i])
int maxProfit_k_inf_with_cool(int[] prices){ int n = prices.length; int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE; int dp_pre_0 = 0; // 代表dp[i-2][0] for (int i = 0; i < n; i++){ int temp = dp_i_0; dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]); dp_i_1 = Math.max(dp_i_1, dp_pre_0 - prices[i]); dp_pre_0 = temp; } return dp_i_0; }
-
k = 任意正整数
一次交易由买入和卖出构成,至少需要两天。所以说有效的限制次数 k 应该不超过 n/2,如果超过,就没有约束作用了,相当于 k = +infinity。这种情况是之前解决过的。
int maxProfit_k_any(int max_k, int[] prices) { int n = prices.length; if (max_k > n/2) return maxProfit_k_inf(prices); // 1 <= k <= max_k int[][][] dp = new int[n][max_k + 1][2]; for(int i = 0; i < n; i++){ for(int k < max_k; k >= 1; k--){ if(i-1 == -1){ //... } // 处理base case // 这里只依赖i-1 的k,所以这里k反向遍历时,所有k的状态都已知(已在外层循环的前一次迭代中计算出来了) dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i]); dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i]); } } return dp[n-1][max_k][0]; }
总结
关键就在于找到所有可能的「状态」,然后想想怎么更新这些「状态」。一般用一个多维 dp 数组储存这些状态,从 base case 开始向后推进,推进到最后的状态,就是我们想要的答案。