固定状态 + 最大增量问题：按策略买卖股票的最佳时机

景真
63 阅读8分钟

问题

给你两个整数数组 prices 和 strategy，其中：

  • prices[i] 表示第 i 天某股票的价格。
  • strategy[i] 表示第 i 天的交易策略，其中：
    • -1 表示买入一单位股票。
    • 0 表示持有股票。
    • 1 表示卖出一单位股票。

同时给你一个 偶数 整数 k，你可以对 strategy 进行 最多一次修改。一次修改包括：

  1. 选择 strategy 中恰好 k 个 连续 元素。
  2. 将前 k / 2 个元素设为 0（持有）。
  3. 将后 k / 2 个元素设为 1（卖出）。

利润定义为所有天数中 strategy[i] * prices[i] 的 总和。

返回你可以获得的 最大 可能利润。

注意：没有预算或股票持有数量的限制，因此所有买入和卖出操作均可行，无需考虑过去的操作。

示例 1：

  • 输入： prices = [4,2,8], strategy = [-1,0,1], k = 2
  • 输出： 10

解释：

修改策略利润计算利润
原始[-1, 0, 1](-1 × 4) + (0 × 2) + (1 × 8) = -4 + 0 + 84
修改 [0, 1][0, 1, 1](0 × 4) + (1 × 2) + (1 × 8) = 0 + 2 + 810
修改 [1, 2][-1, 0, 1](-1 × 4) + (0 × 2) + (1 × 8) = -4 + 0 + 84

因此，最大可能利润是 10，通过修改子数组 [0, 1] 实现。

示例 2：

  • 输入： prices = [5,4,3], strategy = [1,1,0], k = 2
  • 输出： 9

解释：

修改策略利润计算利润
原始[1, 1, 0](1 × 5) + (1 × 4) + (0 × 3) = 5 + 4 + 09
修改 [0, 1][0, 1, 0](0 × 5) + (1 × 4) + (0 × 3) = 0 + 4 + 04
修改 [1, 2][1, 0, 1](1 × 5) + (0 × 4) + (1 × 3) = 5 + 0 + 38

因此，最大可能利润是 9，无需任何修改即可达成。

提示：

  • 2 <= prices.length == strategy.length <= 10^5
  • 1 <= prices[i] <= 10^5
  • -1 <= strategy[i] <= 1
  • 2 <= k <= prices.length
  • k 是偶数

解法

这个问题要求我们找到在最多进行一次特定修改后，能够获得的最大利润。核心在于理解这次“修改”会如何影响总利润，并高效地找出能带来最大利润增长的修改方案。

第一步：理解问题的核心

总利润的计算公式是 总利润 = Σ (strategy[i] * prices[i])，其中 Σ 表示对所有天数 i 求和。

我们的目标是让这个总利润最大化。我们有两种选择：

  1. 不进行任何修改，利润就是原始策略计算出的利润。
  2. 进行一次修改，选择一个起始点 i，修改 strategy 数组中 [i, i + k - 1] 这个长度为 k 的窗口。修改后的新策略会产生一个新的利润。

我们需要在“不修改的利润”和“所有可能修改产生的最大利润”中取一个最大值。

这可以转化为一个更简单的问题：计算出原始的基础利润，然后找出哪一次修改能够带来最大的利润增长量（delta） 。如果这个最大的增长量是正数，我们就执行这次修改；如果是负数或零，我们就不做任何修改（相当于增长量为0）。

所以，最终答案是：最大利润 = 基础利润 + max(0, 最大利润增长量)。

第二步：剖析利润变化量 (Delta)

让我们来分析，当选择从索引 i 开始，长度为 k 的窗口进行修改时，利润会发生怎样的变化。

这个窗口是 [i, i + k - 1]。

  • 前 k/2 个元素 (索引从 i 到 i + k/2 - 1) ：策略被强制修改为 0 (持有)。
  • 后 k/2 个元素 (索引从 i + k/2 到 i + k - 1) ：策略被强制修改为 1 (卖出)。

对于这个窗口内的每一天 j，利润的变化是 (新策略[j] * prices[j]) - (旧策略[j] * prices[j])。

我们把这个窗口分成两部分来分析：

  1. 第一部分 (索引 j 从 i 到 i + k/2 - 1) :

    • 旧的利润贡献: strategy[j] * prices[j]
    • 新的利润贡献: 0 * prices[j] = 0
    • 利润变化量: 0 - (strategy[j] * prices[j]) = -strategy[j] * prices[j]

  2. 第二部分 (索引 j 从 i + k/2 到 i + k - 1) :

    • 旧的利润贡献: strategy[j] * prices[j]
    • 新的利润贡献: 1 * prices[j] = prices[j]
    • 利润变化量: prices[j] - (strategy[j] * prices[j]) = prices[j] * (1 - strategy[j])

因此，修改窗口 [i, i + k - 1] 带来的总利润增长量 delta[i] 是这两部分变化量之和：

δ[i]=∑j=ii+k/2−1(−strategy[j]∗prices[j])+∑j=i+k/2i+k−1(prices[j]∗(1−strategy[j]))\delta[i] = \sum_{j=i}^{i+k/2-1} (-strategy[j] * prices[j]) + \sum_{j=i+k/2}^{i+k-1} (prices[j] * (1 - strategy[j]))

第三步：如何高效地计算所有可能的 Delta

我们现在需要计算所有可能的起始位置 i (从 0 到 n-k) 对应的 delta[i]，然后找到其中的最大值。

如果对每个 i 都进行一次完整的 k 次循环来计算 delta[i]，总的时间复杂度将是 O(n * k)。考虑到 n 和 k 的最大值，这可能会超时。

这里有一个经典的优化技巧：前缀和 (Prefix Sum) 。

我们可以预先计算出两个数组的“变化潜力”：

  1. 一个数组 loss_potential，其中 loss_potential[j] = -strategy[j] * prices[j]。这代表了如果第 j 天的策略变为 0，利润会如何变化。
  2. 一个数组 gain_potential，其中 gain_potential[j] = prices[j] * (1 - strategy[j])。这代表了如果第 j 天的策略变为 1，利润会如何变化。

然后，我们为这两个“潜力”数组分别计算前缀和：

  • prefix_loss[j] = loss_potential[0] + ... + loss_potential[j-1]
  • prefix_gain[j] = gain_potential[0] + ... + gain_potential[j-1]

有了前缀和数组，我们就可以在 O(1) 的时间内计算出任意区间的和：

  • Σ_{j=a}^{b} loss_potential[j] = prefix_loss[b+1] - prefix_loss[a]
  • Σ_{j=a}^{b} gain_potential[j] = prefix_gain[b+1] - prefix_gain[a]

这样，我们计算 delta[i] 的公式就变成了：

delta[i] = (prefix_loss[i+k/2] - prefix_loss[i]) + (prefix_gain[i+k] - prefix_gain[i+k/2])

现在，我们可以通过一次遍历（i 从 0 到 n-k）来计算所有 delta[i]，每次计算都是 O(1)。总的时间复杂度就从 O(n*k) 优化到了 O(n)。

第四步：算法步骤总结

  1. 计算基础利润：遍历整个 prices 和 strategy 数组，计算不进行任何修改时的总利润 baseProfit。

  2. 预计算前缀和：

    • 创建两个大小为 n+1 的前缀和数组 prefixLoss 和 prefixGain，并初始化为0。

    • 遍历 i 从 0 到 n-1：

      • 计算 loss_potential[i] = -strategy[i] * prices[i]。
      • 计算 gain_potential[i] = prices[i] * (1 - strategy[i])。
      • 更新前缀和数组：prefixLoss[i+1] = prefixLoss[i] + loss_potential[i]，prefixGain[i+1] = prefixGain[i] + gain_potential[i]。

  3. 寻找最大利润增长量：

    • 初始化 maxDelta = 0 (因为我们可以选择不修改，所以增长量至少是0)。

    • 令 halfK = k / 2。

    • 遍历所有可能的修改窗口，即 i 从 0 到 n-k：

      • 使用前缀和公式计算当前窗口的 currentDelta。
      • lossPart = prefixLoss[i + halfK] - prefixLoss[i]
      • gainPart = prefixGain[i + k] - prefixGain[i + halfK]
      • currentDelta = lossPart + gainPart
      • 更新 maxDelta = max(maxDelta, currentDelta)。

  4. 计算最终结果：返回 baseProfit + maxDelta。

注意：由于价格和数组长度都可能很大，利润的计算结果可能会超出32位整数的范围，因此需要使用64位长整型（在Java中是 long）来存储利润和前缀和。

Java 代码实现

class Solution {
    public long maxProfit(int[] prices, int[] strategy, int k) {
        int n = prices.length;

        // 步骤 1: 计算基础利润
        long baseProfit = 0;
        for (int i = 0; i < n; i++) {
            baseProfit += (long) strategy[i] * prices[i];
        }

        // 步骤 2: 预计算前缀和
        // prefixLoss[i] 存储的是 [0, i-1] 区间内，策略变为0的利润变化量之和
        long[] prefixLoss = new long[n + 1];
        // prefixGain[i] 存储的是 [0, i-1] 区间内，策略变为1的利润变化量之和
        long[] prefixGain = new long[n + 1];

        for (int i = 0; i < n; i++) {
            // 如果第 i 天策略变为 0 (持有), 利润变化量
            long lossPotential = (long) -strategy[i] * prices[i];
            // 如果第 i 天策略变为 1 (卖出), 利润变化量
            long gainPotential = (long) prices[i] * (1 - strategy[i]);

            prefixLoss[i + 1] = prefixLoss[i] + lossPotential;
            prefixGain[i + 1] = prefixGain[i] + gainPotential;
        }

        // 步骤 3: 寻找最大利润增长量
        long maxDelta = 0; // 如果所有修改都导致利润下降，我们选择不修改，增量为0
        int halfK = k / 2;

        // 遍历所有可能的修改窗口起点 i
        // 窗口范围是 [i, i + k - 1]
        for (int i = 0; i <= n - k; i++) {
            // 第一部分 [i, i + halfK - 1] 策略变为 0
            // 使用前缀和计算这部分的利润总变化
            long lossPart = prefixLoss[i + halfK] - prefixLoss[i];

            // 第二部分 [i + halfK, i + k - 1] 策略变为 1
            // 使用前缀和计算这部分的利润总变化
            long gainPart = prefixGain[i + k] - prefixGain[i + halfK];
            
            long currentDelta = lossPart + gainPart;
            
            if (currentDelta > maxDelta) {
                maxDelta = currentDelta;
            }
        }

        // 步骤 4: 计算最终结果
        return baseProfit + maxDelta;
    }
}

示例 1 演练

我们用 prices = [4,2,8], strategy = [-1,0,1], k = 2 来走一遍流程。

  1. 基础利润: (-1 * 4) + (0 * 2) + (1 * 8) = -4 + 0 + 8 = 4。

  2. 前缀和: n=3, k=2, halfK=1

    • loss_potential = [-(-1)*4, -0*2, -1*8] = [4, 0, -8]
    • gain_potential = [4*(1-(-1)), 2*(1-0), 8*(1-1)] = [8, 2, 0]
    • prefixLoss = [0, 4, 4, -4]
    • prefixGain = [0, 8, 10, 10]

  3. 寻找最大Delta:

    • i = 0: 窗口 [0, 1]。

      • 第一部分 [0,0]: lossPart = prefixLoss[1] - prefixLoss[0] = 4 - 0 = 4。
      • 第二部分 [1,1]: gainPart = prefixGain[2] - prefixGain[1] = 10 - 8 = 2。
      • currentDelta = 4 + 2 = 6。
      • maxDelta 更新为 6。

    • i = 1: 窗口 [1, 2]。

      • 第一部分 [1,1]: lossPart = prefixLoss[2] - prefixLoss[1] = 4 - 4 = 0。
      • 第二部分 [2,2]: gainPart = prefixGain[3] - prefixGain[2] = 10 - 10 = 0。
      • currentDelta = 0 + 0 = 0。
      • maxDelta 仍为 6。

  4. 最终结果: baseProfit + maxDelta = 4 + 6 = 10。这与示例输出一致。

复杂度分析

  • 时间复杂度: O(n)。计算基础利润需要 O(n)，计算前缀和需要 O(n)，遍历窗口寻找最大 delta 需要 O(n-k)，也就是 O(n)。总的来说是线性的。
  • 空间复杂度: O(n)。我们使用了两个大小为 n+1 的数组来存储前缀和。
avatar
文章
阅读
粉丝