青训营X豆包MarsCode 技术训练营:股票市场交易策略优化Python3题解| 豆包MarsCode AI 刷题

71 阅读5分钟

股票市场交易策略优化

问题描述

小R近期表现出色,公司决定以股票的形式给予奖励,并允许他在市场上进行交易以最大化收益。给定一个数组,数组中的第 i 个元素代表第 i 天的股票价格。小R需要设计一个算法来实现最大利润。

股票交易规则如下:

  • 小R可以多次买卖股票,但在买入新的股票前必须卖出之前的股票。
  • 每次卖出股票后存在一天的冷冻期,在冷冻期内小R不能购买股票。

你的任务是帮助小R计算出在遵守交易规则的情况下能够获得的最大利润。

  • stocks: 一个整数列表,表示连续几天内的股票价格。

测试样例

样例1:

输入:stocks = [1, 2]
输出:1

样例2:

输入:stocks = [2, 1]
输出:0

样例3:

输入:stocks = [1, 2, 3, 0, 2]
输出:3

样例4:

输入:stocks = [2, 3, 4, 5, 6, 7]
输出:5

样例5:

输入:stocks = [1, 6, 2, 7, 13, 2, 8]
输出:12

解题思路

  1. 理解问题:我们需要计算在遵守交易规则的情况下能够获得的最大利润。交易规则包括可以多次买卖股票,但在买入新的股票前必须卖出之前的股票,并且每次卖出股票后存在一天的冷冻期。

  2. 动态规划:我们可以使用动态规划来解决这个问题。定义两个状态:

    • hold[i]:表示第 i 天持有股票时的最大利润。
    • unhold[i]:表示第 i 天不持有股票时的最大利润。
  3. 状态转移方程

    • hold[i] 可以从 hold[i-1](前一天持有股票)或者 unhold[i-2] - stocks[i](前两天不持有股票,今天买入)转移过来。
    • unhold[i] 可以从 unhold[i-1](前一天不持有股票)或者 hold[i-1] + stocks[i](前一天持有股票,今天卖出)转移过来。
  4. 初始状态

    • hold[0] 初始化为 -stocks[0](第一天买入股票)。
    • unhold[0] 初始化为 0(第一天不持有股票)。

Python3代码(通过豆包Marscode测试)

def solution(stocks):
    if not stocks:
        return 0
    
    n = len(stocks)
    hold = [0] * n
    unhold = [0] * n
    
    # 初始化状态
    hold[0] = -stocks[0]
    unhold[0] = 0
    
    for i in range(1, n):
        # 更新 hold[i] 和 unhold[i]
        hold[i] = max(hold[i-1], # 保持持有
                      (unhold[i-2] if i >= 2 else 0) - stocks[i]) # 前两天不持有,今天买入
        unhold[i] = max(unhold[i-1], # 保持不持有
                        hold[i-1] + stocks[i]) # 前一天持有,今天卖出
    
    # 最终结果是不持有股票的最大利润
    return unhold[-1]

关键步骤解释

  1. 初始化状态hold[0] 和 unhold[0] 的初始化。
  2. 状态转移:在循环中更新 hold[i] 和 unhold[i],注意处理边界情况(如 i < 2 时 unhold[i-2] 不存在)。
  3. 返回结果:最终结果是 unhold[-1],即最后一天不持有股票的最大利润。

冷冻期对交易的影响

冷冻期意味着在卖出股票后的下一天不能买入股票。因此,在状态转移时,我们需要确保在卖出股票后的下一天不进行买入操作。冷冻期对交易的影响可以通过在状态转移方程中引入一个额外的条件来处理。在更新 hold[i] 时,使用 unhold[i-2] 而不是 unhold[i-1],以确保在卖出股票后的下一天不进行买入操作。hold[i] = max(hold[i-1], (unhold[i-2] if i >= 2 else 0) - stocks[i]):这里使用 unhold[i-2] 来处理冷冻期,确保在卖出股票后的下一天不进行买入操作。

时间复杂度

当前代码的主要时间复杂度来自于两个部分:

  1. 初始化状态:初始化 hold 和 unhold 数组的时间复杂度是 O(n),其中 n 是股票价格数组的长度。
  2. 状态转移:在循环中更新 hold[i] 和 unhold[i] 的时间复杂度是 O(n),因为每个状态的更新都是常数时间操作。

因此,总的时间复杂度是 O(n)

空间复杂度

当前代码的空间复杂度主要来自于两个部分:

  1. 状态数组hold 和 unhold 数组的空间复杂度是 O(n),其中 n 是股票价格数组的长度。
  2. 其他变量:其他变量(如 n)的空间复杂度是 O(1)

因此,总的空间复杂度是 O(n)

代码优化

虽然当前代码的时间复杂度和空间复杂度已经是 O(n),但我们可以通过优化空间复杂度来进一步改进代码。具体来说,我们可以使用滚动数组(rolling array)来减少空间复杂度。

优化后的代码(通过豆包Marscode)

def solution(stocks):
    if not stocks:
        return 0
    
    n = len(stocks)
    
    # 初始化状态
    hold_prev = -stocks[0]
    unhold_prev = 0
    unhold_prev2 = 0
    
    for i in range(1, n):
        # 更新 hold 和 unhold
        hold = max(hold_prev, unhold_prev2 - stocks[i])
        unhold = max(unhold_prev, hold_prev + stocks[i])
        
        # 更新滚动数组
        hold_prev, unhold_prev, unhold_prev2 = hold, unhold, unhold_prev
    
    # 最终结果是不持有股票的最大利润
    return unhold_prev

关键步骤解释

  1. 初始化状态:使用 hold_prevunhold_prev 和 unhold_prev2 来存储前一天和前两天的状态。
  2. 状态转移:在循环中更新 hold 和 unhold,并使用滚动数组的方式更新 hold_prevunhold_prev 和 unhold_prev2
  3. 返回结果:最终结果是 unhold_prev,即最后一天不持有股票的最大利润。

优化后的时间复杂度和空间复杂度

  • 时间复杂度:仍然是 O(n),因为每个状态的更新仍然是常数时间操作。
  • 空间复杂度:优化后的空间复杂度是 O(1),因为我们只使用了常数个额外的变量来存储状态。

通过这种优化,我们可以在保持时间复杂度不变的情况下,显著减少空间复杂度。