LeetCode 3296 移山题：优先队列模拟与二分答案全解析

LeetCode 3296 这道移山题我是在周赛 430 里第一次遇到的。题目读第一遍的时候就隐隐觉得该用优先队列——多人并行、每次分配一个单位工作量、求最少时间这三个信号叠在一起基本就是堆模拟的固定题型。但真正写代码时我才发现优先队列的解法里藏着一个容易想偏的递推关系稍不留神答案就会差出好几倍。这篇文章把这道题的完整解法、为什么贪心是对的、二分答案的平行思路以及我调试时踩过的坑都整理出来。1. 破题思路为什么最少秒数直接指向堆1.1 题目到底在说什么LeetCode 3296 的题面信息量不大但规则里有一个关键限制很容易被忽略。给一个整数mountainHeight表示山的高度一个数组workerTimes表示每个工人的基础耗时。每个工人每次干活可以降低 1 单位高度但第 1 次需要workerTimes[i]秒第 2 次需要2 * workerTimes[i]秒第 3 次需要3 * workerTimes[i]秒以此类推。所有工人同时工作问把山完全移走最少需要多少秒。这里最反直觉的一点是第二次比第一次慢。一个工人连续干活越多后续每次的边际成本就越高。这意味着你不可能让一个速度快的工人包揽全部工作因为他的第二次、第三次、第 k 次会越来越贵最后可能反而不如让另一个工人从第一次开始干。换句话讲这个题的本质是一个调度问题山的高度有mountainHeight份工作要做每份工作必须分配给某个工人的某个工作槽位。我们要给所有单位工作量挑槽位使得最后一个槽位的完成时刻尽可能早。1.2 一条时间线上的工作槽位把每个工人的可工作时间列出来可以看成一个无限长的序列。工人 i 的槽位时间是第 1 次w第 2 次总时间w 2w 3w第 3 次总时间w 2w 3w 6w第 k 次总时间w * (1 2 ... k)也就是w * k * (k 1) / 2注意这里说的总时间是相对于时间零点计算的因为工人从同一时刻开始并行工作。如果我们把所有工人的所有槽位完成时间放在同一条时间线上比如[w1, 2w1, 3w1, ..., w2, 2w2, 3w2, ...]那么想完成mountainHeight次操作只需要取这条时间线上最小的mountainHeight个值即可。为什么这个贪心是严格成立的因为所有操作是并行的整个工程的完工时间只取决于最后一个被选中的槽位在哪一刻完成。如果你不选前mountainHeight个最小槽位而选了某个更大的槽位那么完工时间只会更晚不可能更早。所以问题的答案就是所有工人的所有槽位完成时间合并排序后第mountainHeight个最小值。1.3 优先队列在这里扮演什么角色有了上面的抽象剩下的问题就是如何高效拿到所有槽位时间中第mountainHeight小的值。每个工人的槽位序列都有序递增但多个有序序列合在一起并不是全局有序。这正好是最小堆最擅长的事情把每个工人当前可用的最小槽位放进堆里每次弹出堆顶就代表消耗掉一个槽位然后再把这个工人的下一个槽位补进堆。重复mountainHeight次最后一次弹出的值就是答案。整个过程和合并多条有序链表的思路完全一致。这个算法的时间复杂度是O(mountainHeight log n)空间复杂度O(n)。只要mountainHeight不超过10^5级别跑起来非常轻松。2. 堆模拟的每个细节从状态三元组到递推公式2.1 状态设计绝对时刻与已做次数堆里每个元素需要存三个信息该工人下一次可以完成操作的绝对时刻、该工人的基础耗时w、该工人已经完成的次数cnt。用(nextTime, w, cnt)三元组压入堆。排序时nextTime是第一关键字这样堆顶一定是当前所有可用槽位里结束时间最早的。初始化时每个工人第 1 次操作还没开始已经完成的次数是 0第一次完成时刻正好是w。所以初始入堆工人 i 入堆(workerTimes[i], workerTimes[i], 0)这里有个我后来才注意到的细节如果初始化时把cnt写成 1那么后续所有递推都会错位答案会整体偏大。原因在后面的递推公式里会解释清楚。2.2 核心递推下一次完成时刻怎么算每次操作步骤如下pop堆顶得到(time, w, cnt)其中time是这个工人当前这次操作完成的绝对时刻。因为消耗掉一个槽位让这个工人工作次数加一cnt 1。这个工人现在准备第cnt 1次工作而这一次工作需要耗时w * (cnt 1)。所以新的绝对完成时刻是time w * (cnt 1)。重新把(time w * (cnt 1), w, cnt 1)压入堆。上面写成代码会更直观time, w, cnt heapq.heappop(heap) cnt 1 next_time time w * cnt heapq.heappush(heap, (next_time, w, cnt))我们来验证一下这个递推假设w 2初始cnt 0入堆(2, 2, 0)。第一次 pop 得到time 2cnt 1变成 1新的绝对时刻是2 2 * 1 4这正好是第二次工作2 * 2 4秒从时间 2 到时间 4 的结束点。接下来cnt 2第三次工作耗时2 * 3 6秒下一个绝对时刻是4 4?不对是4 2 * 2 8还是错了我们重新算第二次结束时time 4cnt 1pop 后cnt 1变成 2新时刻是4 w * cnt 4 4 8。但工人第三次工作本身耗时是3 * w 6秒从 4 到 8 只差了 4 秒这明显不对。啊这里我发现了问题。让我重新梳理。关键点next_time应不应该从当前time往后累加这一次的耗时如果工人的第 k 次工作耗时是w*k那么第一次的耗时是w*1第二次的耗时是w*2。初始绝对时刻是 0。第一次工作从 0 开始耗时 w结束在 w。所以第一次完成后绝对时刻 0 w w已做次数 cnt 1。第二次工作耗时 2w结束在 w 2w 3w。此时 cnt 2。第三次耗时 3w结束在 3w 3w 6w。所以在堆状态里如果time表示当前这轮工作完成后的绝对时刻cnt表示已经完成了几次那么下一次工作的耗时是w * (cnt 1)下一次完成时刻是time w * (cnt 1)。初始状态(w, w, 1)。注意这里必须是cnt 1因为已经完成了一次否则按上面公式会多算。刚才我推导过程中把cnt的含义弄混了。正确的写法有两种写法 Acnt 表示已完成次数 初始化(w, w, 1)。pop 后已知time是这次完成时刻cnt是已完成次数。下一次耗时w * (cnt 1)。压入(time w * (cnt 1), w, cnt 1)。写法 Bcnt 表示已完成次数但递增后再用 初始化(w, w, 0)pop 后cnt 1此时cnt变成该工人这次刚完成的次数。下一次要干的是第cnt 1次工作耗时w * (cnt 1)。压入(time w * (cnt 1), w, cnt 1)。两种写法拿到的 next_time 是一样的。我 2.2 里最初那段推导混合了两者写了cnt 1后又直接用了time w * cnt这相当于把当前这次的耗时叠加了变成从旧 time 到新 time 只增加w * cnt自然是错的。错误点在于当cnt递增到 k 后正确的下次耗时应该是w * k这个数本身而不是w * (k1)再加一次。也就是说next_time time w * (cnt 1)而不是time w * cnt。我把正确的递推在 2.3 代码里写清楚并在第 5 节单独强调这个坑。2.3 可 AC 的 Python 实现import heapq from typing import List class Solution: def minimumSeconds(self, mountainHeight: int, workerTimes: List[int]) - int: heap [] for w in workerTimes: # (下一次完成绝对时刻, 基础耗时, 该工人已完成次数) heapq.heappush(heap, (w, w, 1)) ans 0 for _ in range(mountainHeight): time, w, cnt heapq.heappop(heap) ans time # 第 cnt 1 次工作耗时 w * (cnt 1) next_time time w * (cnt 1) heapq.heappush(heap, (next_time, w, cnt 1)) return ans这段代码里每次pop出来的time都是当前全局最早可用的槽位完成时间。循环进行mountainHeight次后最后一次time就是我们需要的第mountainHeight小的槽位时间。如果想压一点常数可以把初始化改成heap [(w, w, 1) for w in workerTimes] heapq.heapify(heap)这里有个容易被忽略的小优化heapify比逐push快尤其是在n很大的时候。2.4 时间和空间复杂度堆里始终有n个元素每次pop和push都是O(log n)。外层循环执行mountainHeight次所以总时间复杂度是O(mountainHeight * log n)。空间上只需要维护一个长度为n的堆所以空间复杂度是O(n)。当mountainHeight比较小时比如10^5以内这个解法又快又稳。但一旦mountainHeight涨到10^7甚至10^9模拟mountainHeight次就完全不现实了这时候得换二分答案的思路。3. 二分答案的平行解法什么时候该换思路3.1 二分框架与单调性优先队列模拟是逐单位时间推进二分答案则是把时间当作自变量直接验证某个时间点是否足够。单调性很明显如果给定T秒内所有工人能把山移完那么任何大于T的时间也一定可以如果T秒内移不完那么任何小于T的时间也移不完。因此答案落在[0, 上界]区间内可以对时间做二分。重点在于如何写判定函数check(T)给定T秒每个工人最多能完成多少次操作把所有工人能完成的次数加起来如果大于等于mountainHeight说明T可行。3.2 单个工人的能力上限解二次不等式对于基础耗时为w的工人如果他在T秒内完成了k次操作必须满足w * (1 2 ... k) T左边累加得到w * k * (k 1) / 2 T也就是k * (k 1) 2 * T / w这是个一元二次不等式。最稳妥的做法是在判定函数内部再做一个整数二分找最大的k满足条件。二分左边界是 0右边界可以设为mountainHeight因为一个人不可能干超过总山高的次数。def can(T: int) - bool: total 0 for w in workerTimes: lo, hi 0, mountainHeight while lo hi: mid (lo hi 1) // 2 if w * mid * (mid 1) // 2 T: lo mid else: hi mid - 1 total lo if total mountainHeight: return True return False这段代码最需要注意的点是w * mid * (mid 1) // 2可能溢出int。Python 的整数是无限精度所以没事如果改成 C建议直接long long并且先用mid * (mid 1)除以 2 再乘w依然可能溢出时就用__int128。其实更高效的方式是用math.isqrt直接估算根再往下修正。因为要解k^2 k - 2T/w 0取正根约等于k ≈ (isqrt(1 4 * (2 * T // w)) - 1) // 2然后微调一两次。这个做法能省掉一层二分但容易出精度边界问题我一般还是用整数套整数二分写起来放心。3.3 两种解法的对比与选型对比项优先队列模拟二分答案核心思想事件驱动每次取最快槽位把时间作为自变量验证可行性时间复杂度O(mountainHeight log n)O(n * log(上界) * log(mountainHeight))对山高规模的要求需要mountainHeight可枚举完全不受mountainHeight大小影响代码出错点堆状态递推容易写反判定函数二次不等式和溢出面试/周赛推荐度mountainHeight 1e5时首选作为优化方案展示时加分选型建议很直接如果mountainHeight在10^5以内优先队列解法在思路和代码上都更贴近题意基本不容易错。只有在mountainHeight大到无法枚举时才必须转二分。实际竞赛里也有一些出题人故意把二者结合比如先用二分缩小范围再堆模拟但那是极少数情况。3.4 二分代码的一点补充二分时间上界最稳的取法是让耗时最小的工人连续把整个山搬完。假设最小基础耗时为min_w那么上界公式是hi min(workerTimes) * mountainHeight * (mountainHeight 1) // 2这个值保证一定可行因为把全部工作交给最快工人一个人完成这是一个合法方案。注意乘法可能很大用 Python 不用管用 C/Java 记得long。二分主循环lo 0 hi min_w * mountainHeight * (mountainHeight 1) // 2 while lo hi: mid (lo hi) // 2 if can(mid): hi mid else: lo mid 1 return lo4. 从递增工时模型看竞赛套路4.1 这类题的特征信号LeetCode 3296 其实属于一类特征非常明显的递增工时模型。识别这类题的三个信号有多个执行者或者多个并行通道目标是把某个总量消耗完比如把山移平、把任务做完、把香蕉吃完执行者的单次成本随次数递增比如第 k 次成本是k * w、2^k或者k^2。一旦同时出现这三个特征基本可以确定两条路优先队列逐次选最优或者二分总量/时间验证可行性。3296 正好把两条路都打通了所以是一道很好的综合练习题。4.2 与 875 爱吃香蕉的狒狒的异同LeetCode 875 是经典的爱吃香蕉的狒狒二分题。大意是给定几堆香蕉和一个时间上限求最少每小时吃多少根香蕉能在时限内吃完。875 和 3296 的共同点是都有一个判定函数给定一个变量判断能否在限制内完成任务。875 的判定函数是对每一堆香蕉做一次除法累加3296 的判定函数则需要解二次不等式因为工人的成本递增。它们的不同点在于875 只有一个执行者狒狒自己不存在多执行者并行调度的问题3296 有多工人并行因此才有优先队列解法存在的空间。如果刷题时想串起来理解可以把 875 当成二分答案入门题把 3296 当成二分答案进阶 堆模拟综合题。两道做完对答案单调性和判定函数设计就会有立体的感觉。4.3 同思路的扩展题清单类似题目可以按两类整理堆模拟类2462 Total Cost to Hire K Workers多指针加堆选前 k 小23 Merge k Sorted Lists合并 k 个有序链表373 Find K Pairs with Smallest Sums多路归并。二分答案类1011 Capacity To Ship Packages Within D Days二分容量774 Minimize Max Distance to Gas Station二分最大间隔2064 Minimized Maximum of Products Distributed to Any Store二分最大值。3296 的位置比较特殊它既是多路有序序列取前 k 小的堆模拟题又是一个验证总量可行性的二分题。把这两套思路在同一个题上都练一遍价值比单纯刷十几道单一类型的题要大。5. 我实际提交时踩过的坑5.1 答案规模和整数类型第一眼看到这道题的返回值可能会觉得秒数嘛int 就行了。但实际上答案很容易突破2^31 - 1。最简单的例子workerTimes [1]mountainHeight 100000。只有一个工人他完成所有工作的总时长是1 * 100000 * 100001 / 2 5000050000秒。这个数已经超过 32 位整数上限更不要说workerTimes[i]可以到10^6级别。所以在 C 里全部用long longJava 用longPython 不用担心但代码习惯上也不要随意转int。5.2 堆初始化和递推的经典错误我在第 2.2 节推导时差点犯的错误就是弄混cnt的含义。这里直接给结论如果你用三元组(nextTime, w, cnt)表示该工人当前这次工作完成的绝对时刻是 nextTime他已经完成了 cnt 次那么初始化必须是(w, w, 1)不是(w, w, 0)。因为第一次完成后的时刻是w已完成次数是 1。下一次工作时耗时是w * 2所以新的 nextTime 是w w * 2 3w这正好是第二次完成的总用时。如果把cnt初始化为 0按next_time time w * (cnt 1)计算第一次完成时刻time wcnt 0下一次就是w w * 1 2w这对应的是第二次工作耗时 w的错误结果整体答案会偏小而且后面全错。我的建议是先用小数据手算一遍模拟过程确认堆状态的变化再写代码。比如workerTimes [1]mountainHeight 3正确过程是初始(1, 1, 1)pop 得到 1完成第一次压入(1 1*2, 1, 2) (3, 1, 2)pop 得到 3完成第二次压入(3 1*3, 1, 3) (6, 1, 3)pop 得到 6完成第三次答案 6手动跑一遍哈希值和公式对不对一目了然。5.3 二分上界与溢出二分答案的上界如果随便设一个很大的数比如10^18在 C 里mid * (mid 1)可能溢出long long进而导致判定函数出错甚至死循环。最稳妥的上界就是最快工人单干完全部活的时间hi min(workerTimes) * mountainHeight * (mountainHeight 1) // 2这个值是合法解且足够紧不容易溢出。在 C 中如果担心乘法中间结果溢出可以这样写long long hi minW * mountainHeight * (mountainHeight 1) / 2;因为minW、mountainHeight都是long long乘法从左到右运算只要总数不超过9e18就是安全的。一旦题目把mountainHeight抬到10^9这个上界会到5e23就必须用__int128或者换一种上界策略。5.4 建议的测试用例我实际调试时会固定跑这几个用例mountainHeight 1, workerTimes [5]答案5验证初始状态正确。mountainHeight 3, workerTimes [1, 1]。两个工人的槽位完成时间分别是1, 3, 6, ...和1, 3, 6, ...。整体排序前三个是1, 1, 3答案应该是3。这个用例能检验并行是否正确两个工人在第 1 秒同时完成第一次随后一人完成第二次的时刻是第 3 秒所以 3 秒完工。mountainHeight 3, workerTimes [1, 2]。槽位时间工人 A 是1, 3, 6...工人 B 是2, 6, 12...。前三个最小槽位是1, 2, 3答案3。这个用例能检验堆是否会在不同工人之间正确切换。mountainHeight 100000, workerTimes [1]。答案5000050000用来验证整数溢出。把这些用例跑完优先队列解法和二分解法基本都能一次过。我个人在实际操作中的体会是这类递增工时的题堆模拟写着爽但真正决定代码对错的往往是状态定义那一行建议先把cnt的语义在注释里写清楚再开始敲代码。