力扣75颜色分类:一次遍历的原地三指针解法与循环不变式

📅 发布时间:2026/10/13 10:39:42
力扣75颜色分类:一次遍历的原地三指针解法与循环不变式
如果你刷到力扣hot100里的颜色分类大概率会在评论区看到同一个感慨“这题看着比两数之和还简单一动手全是细节。”确实题目描述一句话就能说完——给定一个只含0、1、2的数组把它们原地排好序——可当你试着自己实现一次遍历版的AC代码时指针的进退、交换的时机、循环的边界每一个地方都能埋雷。今天这篇就是来拆雷的顺便把这道经典题背后的partition思想、循环不变式、以及在快排优化里的应用都串一遍看完你不仅能AC还能在面试时讲得比面试官期待得还细。1. 这道题到底在考什么不只是“把012排个序”1.1 题目的真实出处与约束条件颜色分类在力扣上的题号是75很多人叫它荷兰国旗问题。这个说法来自一篇非常经典的计算机科学论文当年论文用红白蓝三色国旗做类比描述的就是把一个只含三种值的序列排成三段的问题。力扣把它改写成0、1、2之后难度标的是中等但说实话它的算法思想并不难难的是你能不能在一堆细节里不出错。题目给的数组里元素只可能是3种值传统排序算法面对的是任意可比较元素而这题只需要利用“只有3种值”这个特性。同时题目要求原地修改不能用额外数组。力扣上这道题的Follow Up其实写得很直白一个直接思路是两次遍历的计数排序但你能不能在一次遍历、常数空间的约束下完成这两个约束就是这道题的灵魂。如果只是把数组排序API调一行就结束了如果允许额外数组计数排序一遍统计两遍回填也就结束了。正因为加上了“一次遍历、原地、常数空间”你才被迫去思考双指针甚至三指针的方案也正因为如此这道题才值得被收进hot100。1.2 为什么hot100收录它从数据结构到partition思维我在准备面试的过程中发现hot100里很多题目其实不是考你会不会某种高级算法而是考你有没有形成“在约束下设计算法”的肌肉记忆。颜色分类是一个非常典型的代表它表面上是一个数组排序问题实际上考察的是——你知不知道快速排序里的partition操作你能不能把“排序”这个抽象目标转化成一个具体的循环不变式你写出来的循环能否用数学归纳法式的推理说服自己和面试官。如果只看数据规模O(n log n)的排序算法当然能过但一个合格的候选人在看到“数组中只有0、1、2”时第一反应应该是“这里存在比通用排序更快的办法”。这种条件意识正是工程里优化代码时最常用的思维方式——看到数据有特殊分布就要思考能不能用更廉价的方案。另外一个隐藏考点是边界处理。三指针算法代码不超过15行但新手写出来往往在某个测试用例上翻车翻车点集中在那几个if分支的写法和循环条件上。能准确解释“为什么遇到2交换后i不能立即后移”的人才是真正理解了这个算法而不是背代码。2. 两种“能过但别急着说”的常规解法sort与计数2.1 直接排序能AC但暴露了一个思维定式先看最直接的解法def sortColors(nums): nums.sort()就一行。力扣里这道题你用Python提交这行代码能通过因为题目只要求“把数组排好序”没有任何地方禁止你调用排序API。但作为准备面试的人我建议你不要一上来就说这个答案原因有两点。第一调用sort确实完成了任务但完全浪费了“元素只有3种”这个条件。面试官接下来的问题几乎是固定的“如果数组长度是一亿呢如果这是一个嵌入式环境内存只有几KB呢”你当然可以说“通用排序也很稳定”但面试官想要的是一个能根据输入特性定制方案的人。第二sort的复杂度是O(n log n)而这道题用计数排序能轻松做到O(n)。一个O(n)能解决的问题你用O(n log n)面试官不会说你错但会怀疑你的敏感度。在真实工程中这种“用重型工具解决轻量问题”的做法往往意味着无谓的资源浪费。我个人的建议是这道题可以当作“先讲低配再讲高配”的典型题目来练。你甚至在面试里可以主动说“如果允许两遍遍历我会用计数排序但如果要求在单遍遍历和常数空间下完成我会用三指针。”这种主动递进式的回答本身就展示了你的思考层次。2.2 计数排序O(n)时间但不是单遍遍历这是很多人的第二反应也确实是个好答案def sortColors(nums): cnt [0, 0, 0] for x in nums: cnt[x] 1 idx 0 for val in range(3): for _ in range(cnt[val]): nums[idx] val idx 1原理很直接第一遍统计每个颜色出现了几次第二遍从头回填。因为颜色只有3种cnt数组的长度是固定的常数3所以额外空间是O(1)时间也是O(n)。这个解法对于“把一堆对象按固定类别分组”的场景非常实用工程上也经常用类似思路做桶统计。但请注意它是“两遍遍历”的算法。力扣的Follow Up明确问你“能不能在一次遍历内完成”这才是这道题真正的进阶点。每次想到计数排序你得先确认需求的约束是什么——如果产品经理只要求功能实现两次遍历完全够用如果这是一个对延迟极度敏感的在线服务数据进来只能扫一遍那就得用单遍方案。2.3 面试官真正想听的是什么有朋友问过我“既然两道常规解法都能过为什么还非得学三指针”我的回答是不是说三指针一定比计数排序好而是你要具备在同一道题上给出多套方案并分析取舍的能力。计数排序适合“类别数量少、统计阶段允许重放数据”的场景三指针适合“只能一次遍历、必须在流动过程中就地分拣”的场景。面试官真正想考察的是你能否根据约束切换方案。如果你只会sort说明你缺乏对输入特征的敏感如果你会计数排序但不会三指针说明你在进阶约束下缺少进一步优化的手段如果你能写出三指针并且讲清楚循环不变式那这道题基本就满分了。3. 双指针解法左端拦0、右端堵2游标一次扫过3.1 核心思路与代码既然要求一次遍历、原地修改、常数空间那就得想办法在扫描的过程中每遇到一个元素就把它放到最终应该待的位置。怎么知道一个元素最终应该待在哪对于0它一定在数组最左边对于2它一定在数组最右边1自然落在中间。所以我们用三个变量维护边界p0表示下一个0应该放的位置初始是0p2表示下一个2应该放的位置初始是n-1i是当前遍历游标初始是0。算法框架是i从左往右走遇到0就跟p0位置的元素交换、p0右移、i继续遇到2就跟p2位置的元素交换、p2左移、i原地不动遇到1直接i右移。def sortColors(nums): n len(nums) p0, p2, i 0, n - 1, 0 while i p2: if nums[i] 0: nums[i], nums[p0] nums[p0], nums[i] p0 1 i 1 elif nums[i] 2: nums[i], nums[p2] nums[p2], nums[i] p2 - 1 else: i 1这个版本简洁无误我建议你把这个写法当作标准答案记下来。但记下来之前你得先理解它为什么这么写。3.2 手推一遍[2,0,2,1,1,0]每一步指针都在干嘛文字解释总有点虚我们一步步推一个具体例子。初始数组是[2, 0, 2, 1, 1, 0]p00p25i0。步骤当前inums[i]操作操作后数组p0p2i102交换nums[0]与nums[5][0,0,2,1,1,2]040200交换nums[0]与nums[0][0,0,2,1,1,2]141310交换nums[1]与nums[1][0,0,2,1,1,2]242422交换nums[2]与nums[4][0,0,1,1,2,2]232521i右移[0,0,1,1,2,2]233631i右移[0,0,1,1,2,2]234第6步结束之后i4p23循环条件i p2不成立退出。最终数组是[0,0,1,1,2,2]完全正确。细看这个推演会发现在第2步和第3步交换的位置是元素自己跟自己其实是原地没动的交换。这是正常的因为p0和i刚好重合时说明当前位置本来就应该放0。不要觉得这是bug。3.3 交换背后的关键判断什么时候i可以动什么时候不能动这是整个算法的核心也是新手最容易写错的地方。遇到0时nums[i]与nums[p0]交换p0右移i右移。为什么i可以放心右移因为nums[p0]这个位置只可能是0或1不可能2。这里需要一点推理所有2在游标经过它们时都会通过“遇到2”分支被换到p2右侧所以p0到i这个区间里不可能还住着一个2。既然交换过来的值不是2那么要么是0要么是1这两种值都不会影响后续正确性——如果换了0过来那恰好是当前位置应该处理的0i也不会漏如果换了1过来它本来就该待在中间区。遇到2时操作是交换nums[i]与nums[p2]然后p2左移但i不移动。为什么i不能移动因为从p2位置换过来的元素我们完全不知道它是0、1还是2。它是0的话下一步还得把它换到左边它是1的话下一步i可以跳过它是2的话还得继续跟新的p2交换。这就好比你在传送带上看到一件不合格的次品你把它扔到废品区但传送带上刚补上来的新零件质量咋样你必须再看一眼。从p2换过来的数是2这种情况并没有被忽略因为i不动就会再次进入循环重新判断这个位置。这也是为什么循环条件必须写i p2而不是i p2——当i和p2相遇时这个位置还没被处理过不能直接跳过。4. 循环不变式视角为什么这个写法“一定对”4.1 三个区间和一个恒成立的性质代码对了但很多教程不会告诉你它为什么对。这里我用循环不变式来讲这也是面试时你的加分点。每次while循环开始前数组可以被划分成四个区间从0位置到n-1位置[0, p0) 区间内全部是0[p0, i) 区间内全部是1注意p0 i始终成立[i, p2] 是尚未处理的区间(p2, n-1] 区间内全部是2这个划分非常漂亮左边堆0中间堆1右边堆2中间夹着未处理区。初始时p00p2n-1i0四个区间分别是空、空、整个数组、空恒成立。我们只要能证明每次循环操作之后这个不变式依然成立那么循环结束时i p2未处理区间变成空集整个数组就是0、1、2各归其位的结果。4.2 三种分支为什么都能保持不变式分情况来证明。当nums[i] 0时说明未处理区第一个元素是0。把nums[i]与nums[p0]交换此时要分两种细况如果p0 i交换没意义直接把p0和i都右移[0, p0)多了个0保持成立如果p0 i根据不变式nums[p0]属于[p0, i)区间它一定是1交换后nums[i]变成1nums[p0]变成0这时[0, p01)全是0[p01, i]全是1然后p0、i区间定义恢复原样成立。当nums[i] 1时什么都不用做直接i右移。[0, p0)依然全0(p2, n-1]依然全2而[p0, i)的末尾多了一个1仍然全1成立。当nums[i] 2时把nums[i]与nums[p2]交换。交换后nums[p2]变成2这个位置就进入了(p2, n-1]区间p2--之后新的(p2, n-1]全是2成立。但nums[i]变成的是原来nums[p2]的值可能0、1、2它仍然在未处理区[i, p2]里所以i绝对不能动下一次循环再根据新值走到对应分支。循环条件i p2保证了这个未处理区的最后一个元素也能进入循环体不会漏。有了这个证明代码的正确性就不是“看起来对”而是数学上确定的对。这就是为什么我强烈建议你在面试里把这段讲出来——绝大多数候选人能默写代码但只有少数人能讲清楚不变式。4.3 循环边界i p2的由来漏掉最后一个2的真实案例我在给朋友讲这道题时他写的是while i p2然后发现[2,0,1]这个用例排序后变成了[1,0,2]。问题出在哪推演一下就清楚了。nums[2,0,1]p00p22i0。i0nums[0]2交换nums[0]和nums[2]数组变成[1,0,2]p21i不变。i0nums[0]1i变成1。此时i1p21循环条件i p2不成立退出。数组成了[1,0,2]中间的0根本没被处理。如果把条件改成i p2i1时nums[1]0会触发0分支交换后得[0,1,2]正确。为什么会有这个差异因为i p2把i p2时尚未处理的元素直接当成了“已经处理完”。注意p2位置虽然右侧都是2但p2本身可能还不是一个合法的2它可能是0或者1必须让游标i扫到它。这个反例建议你好好记下来面试时如果被问到“为什么不是i p2”直接用[2,0,1]三行推演搞定。提示循环条件里的等于号不能省。i和p2相遇时说明这个位置还没有被处理过必须再进一次循环体。5. 另一条推导路径覆盖赋值法与其他变体5.1 覆盖赋值法先当作2写再回填1和0三指针不是唯一的一次遍历方案。我在刷题时见过一个非常巧妙的覆盖赋值法虽然不太直观但代码短、思想特别适合在面试中展示知识广度。def sortColors(nums): n0, n1 0, 0 for i in range(len(nums)): cur nums[i] nums[i] 2 if cur 2: nums[n1] 1 n1 1 if cur 1: nums[n0] 0 n0 1这个思路是遍历时先把当前格子写成2等于是默认这个位置最终会被一个2占用。然后如果原始值比2小说明它其实该放到1区就在n1位置补一个1n1右移如果原始值比0还小说明它其实是0就在n0位置补一个0n0右移。为什么这样不会乱因为n1和n0始终指向“下一个1”和“下一个0”应该放的位置而2的写入是直接铺底的。换句话说这个算法先把所有的坑都填上2再从后往前把本该是1和0的坑覆盖掉。它在逻辑上等价于“预先分配好所有2再往前插入1和0”因为遍历是从左到右的所以覆盖顺序也不会出错。我用同样的例子验证过一遍结果是正确的。但这种写法的风险在于可读性它不像三指针那样有一目了然的分支逻辑面试让你讲思路时你反而要用更多时间解释两步if。所以我更推荐把覆盖赋值法当作“备胎”——它证明了同一个问题可以有多条推导路径但日常为了AC和讲解三指针仍然是首选。5.2 两次单指针写法更适合解释给队友听的版本还有一种思路是把单指针简化版做成“先放0再放2”的两段式。第一段遍历把所有0换到左边第二段从右往左遍历把所有2换到右边。这样两端一夹中间自然全是1。def sortColors(nums): p0 0 for i in range(len(nums)): if nums[i] 0: nums[i], nums[p0] nums[p0], nums[i] p0 1 p2 len(nums) - 1 for i in range(len(nums) - 1, -1, -1): if nums[i] 2: nums[i], nums[p2] nums[p2], nums[i] p2 - 1这个版本的时间是O(n)严格说两次O(n)常数空间代码更容易理解而且不像三指针那样容易在边界上出错。缺点是它遍历了两次不满足“单次遍历”的进阶要求。但在面试沟通时你可以先讲这个版本再讲三指针让人看到你是如何在“放宽约束”与“加强约束”之间切换的。两次单指针写法对初学者特别友好因为它把“把0往左放”和“把2往右放”两个动作彻底分开了每一步都不用考虑另一个指针的影响。5.3 复杂度与适用场景对比用一个表格总结上面提到的方案方案时间复杂度空间复杂度遍历次数可读性适用场景直接排序O(n log n)O(1)或O(log n)不确定极高无约束时的简单实现计数排序O(n)O(1)颜色固定3种2高类别少、可重放数据两次单指针O(n)O(1)2高想快速解释思路三指针O(n)O(1)1中高必须单遍遍历覆盖赋值O(n)O(1)1中想展示不同思路这张表也是面试时的好材料。你不需要背表格但心里要有这张权衡图。不同约束条件对应不同解法才是算法面试里真正的考点。6. 测试用例设计与易翻车点盘点6.1 一组值得单测的边界输入不管你在力扣上刷了多少题提交前都应该先在本地跑一遍边界用例。颜色分类的测试我建议至少包含以下几类空数组[]任何正确解法都该直接返回。单元素[0]、[1]、[2]交换逻辑不能把同一个位置越界。全相同[0,0,0]、[1,1,1]、[2,2,2]指针移动不能失控。2开头[2,0,1]、[2,2,0,1,1]。0结尾[1,2,0]、[2,1,0]。已经有序[0,1,2]、[0,0,1,1,2,2]。完全逆序[2,2,1,1,0,0]。混合乱序[1,0,2,1,0,2,1,0,2]。我自己习惯把这些用例组织成一个小的测试函数方便一键跑。刷过几轮之后你会发现80%的翻车都集中在“2开头”和“2结尾”这两类用例上因为它们最容易触发p2的交换逻辑和循环边界的临界状态。6.2 刷题过程中的高频bug复盘第一类bug是循环条件。把while i p2写成while i p2遇到2结尾或2与中间值相邻时就会漏处理。前面[2,0,1]的例子就是反例。第二类bug是p2操作后i误自增。有人会在nums[i] 2分支里顺手写个i 1导致从右侧换回来的元素被跳过。比如数组[2,0,1]如果交换后立刻ii1时就跳过了原本p2交换过来的0结果同样错误。记住只有确定当前位置已经处理完i才能向右移动。第三类bug是边界指针的初始值。p2如果写成len(nums)而不是len(nums)-1交换时就会取到数组末尾外的位置直接越界。另一个隐蔽版本是循环条件里用i len(nums)同样会越界。这个错误听起来低级但在紧张状态下很容易发生尤其是从“右边界”的直觉出发时总想给数组末尾留个空位。第四类bug是交换顺序。Python的a,bb,a看起来一行搞定但它会先取右侧两个值再赋值不会出错但如果你改用tmp变量时写错顺序比如tmpnums[i]; nums[i]nums[p2]; nums[p2]tmp一旦写成先修改nums[i]再赋值就会丢数据。建议在纸上推演一次把三个变量的传递顺序画出来。6.3 验证脚本与随机化测试光靠手推用例不够我推荐做一个随机验证脚本让算法跟内置排序对比跑几千次。比如import random def sortColors(nums): n len(nums) p0, p2, i 0, n - 1, 0 while i p2: if nums[i] 0: nums[i], nums[p0] nums[p0], nums[i] p0 1 i 1 elif nums[i] 2: nums[i], nums[p2] nums[p2], nums[i] p2 - 1 else: i 1 for _ in range(5000): arr [random.choice([0, 1, 2]) for _ in range(random.randint(0, 30))] expected sorted(arr) sortColors(arr) if arr ! expected: print(FAIL:, arr, expected) break else: print(ALL PASS)如果代码有隐藏bug随机测试大概率能抓到。这也是我在实际刷题中养成的习惯任何数组排序类题都写一个与内置排序对拍的脚本。它比手推十个固定用例更可靠尤其适合这种边界情况较多的算法。随机测试能覆盖到人脑很难枚举出来的组合比如[2,1,2,0,1,0,2]这种看似乱来的序列可能恰好暴露某个分支的错误。7. 一道题的衍生价值从荷兰国旗到三向切分快排7.1 三向切分快排大量重复键场景下的救星这道题的祖师爷——荷兰国旗问题最初就是研究“如何对大量重复键进行快速排序”时提出来的。经典快速排序在遇到大量相等元素时有一个著名的问题数组里如果有大量重复值每次partition都把这些相等的元素在左半和右半来回搬效率很差最坏能退化到O(n^2)。三向切分快速排序的思想就是借鉴这道题的三指针把数组分成小于基准、等于基准、大于基准三个区间。这样一次partition不仅排好了基准的位置还把所有和基准相等的元素一次性放到了正确区域下一次递归只需要处理小于区和大于区。对于键值重复率高的数据比如数据库里按性别、状态、年龄段分组的记录这种排序方式能省下非常可观的无效比较。我在实际工程中遇到过类似场景上游日志数据里某个状态字段只有3到5种取值却占了上亿条记录。如果对全量记录做通用排序排序时间很难看但用三向切分思路把字段值当作分组键做分区再在各分区内部排序整体耗时能减少一个量级。当然工程上你还得考虑稳定性、内存访问局部性等因素但三向分区的思想是这些优化的共同底座。7.2 同思路的力扣变式移动零、按奇偶排序、三向partition模板颜色分类的思路可以直接平移到很多数组处理的题目上。比如移动零它其实是只有两类元素的特例把0移到左边非0移到右边。一个p0指针就够了def moveZeroes(nums): p0 0 for i in range(len(nums)): if nums[i] ! 0: nums[i], nums[p0] nums[p0], nums[i] p0 1再比如按奇偶排序要求把所有偶数排到奇数前面同样只需要一个左指针本质就是一次partition。如果你愿意再抽象一层可以写出一个通用的三向partition模板给定一个枢轴值pivot把数组分成“小于pivot、等于pivot、大于pivot”三部分。颜色分类其实就是pivot1时把小于1的0放左边、大于1的2放右边。这个通用模板的价值在于面试里出现“把数组按某个阈值分成三段”的题目时你可以直接套用。类似的还有快速选择算法里的三路划分用来在存在大量重复元素时高效找第k大。说白了颜色分类不只是hot100里的一道题它是一类“按值分桶”问题的标准骨架。7.3 我的几点实战体会最后说几句刷题之外的话。颜色分类这道题是我在备面时反复用来练习“讲题”的素材。一个好的讲题节奏是这样的先用一句话说清楚题意然后给出计数排序这种最自然的两遍方案再引出三指针的单遍方案最后用循环不变式证明正确性并用一两个反例强调边界条件。整个过程控制在三分钟以内面试官基本能确认你真正理解了题目。我自己的一个心得体会是与其背三指针代码不如每次刷这类题时先写循环不变式再回到代码里确认每个分支是否保持了不变式。这个方法能治“看起来会做、一写就错”的毛病尤其是对边界条件多的题目特别管用。如果你还在hot100早期阶段建议把这道题列入第一批精刷名单。它不仅是一道面试题更是一把理解partition思想的钥匙。把这把钥匙用好后面碰到快速排序、快速选择、数组分块处理一类的问题你会比别人多一个思考维度。