带k股持仓/负债约束的股票买卖最大收益：求O(n log k)级解法

问题回顾 #

给定n天股票价格序列a₁, a₂, ..., aₙ，以及整数1 ≤ k ≤ n，每天可执行以下三种操作之一：

• 不进行任何操作
• 卖出1股（负债不超过k股即可，持仓为0时也可卖出）
• 买入1股
  约束条件：任何时刻持仓最多k股，或负债最多k股（持仓为负的绝对值不超过k），且第n天后持仓必须为0。已实现O(nk)复杂度的动态规划解法，现寻求更高效的O(n log k)级解法。

核心思路 #

这个问题本质等价于最多完成k次“低买高卖”的配对交易，支持先卖后买（做空），最终所有仓位平仓。我们可以利用两个优先队列（堆）来快速捕捉最优交易机会，通过堆的O(log k)级极值查找操作，将整体复杂度压缩到O(n log k)。

具体实现方案 #

数据结构选择 #

• 最小堆：记录已买入的价格（或做空后的平仓成本），用于快速找到当前最适合卖出的低价持仓。
• 最大堆：记录已卖出的价格（或做空的收益基准），用于快速找到当前最适合平仓的高价空仓。

遍历处理逻辑 #

遍历每天的价格p，依次执行以下步骤：

1. 锁定卖出收益：如果最小堆不为空且当前价格p大于堆顶的低价，卖出该持仓，累计收益p - 堆顶值，同时将p加入最大堆（记录一个反向交易的可能，用于后续调整），弹出堆顶值。
2. 锁定平仓收益：如果最大堆不为空且当前价格p小于堆顶的高价，买入平仓该空仓，累计收益堆顶值 - p，同时将p加入最小堆，弹出堆顶值。
3. 记录潜在机会：将当前价格p加入最小堆，作为后续可能的买入/平仓成本。

复杂度分析 #

每个价格最多入堆、出堆各两次，堆的单次操作复杂度为O(log k)，因此整体时间复杂度为O(n log k)，空间复杂度为O(k)。

伪代码实现 #

import heapq

def max_stock_profit(prices, k):
    min_heap = []  # 存买入成本/平仓成本
    max_heap = []  # 用负号模拟最大堆，存卖出价格/做空基准
    total_profit = 0
    
    for p in prices:
        # 卖出之前的低价持仓
        if min_heap and p > min_heap[0]:
            total_profit += p - heapq.heappop(min_heap)
            heapq.heappush(max_heap, -p)
        # 平仓之前的高价空仓
        if max_heap and p < -max_heap[0]:
            total_profit += (-max_heap[0]) - p
            heapq.heappop(max_heap)
            heapq.heappush(min_heap, p)
        # 加入当前价格作为潜在交易机会
        heapq.heappush(min_heap, p)
    
    return total_profit

注：Python原生仅支持最小堆，因此用负号将最大堆转换为最小堆处理，逻辑完全等价。

内容的提问来源于stack exchange，提问作者MangoPizza