题目:
https://leetcode-cn.com/problems/sort-colors/
给定一个包含红色、白色和蓝色,一共 n 个元素的数组,原地对它们进行排序,使得相同颜色的元素相邻,并按照红色、白色、蓝色顺序排列。
此题中,我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。
注意:
不能使用代码库中的排序函数来解决这道题。
示例:
输入: [2,0,2,1,1,0]
输出: [0,0,1,1,2,2]
进阶:
一个直观的解决方案是使用计数排序的两趟扫描算法。
首先,迭代计算出0、1 和 2 元素的个数,然后按照0、1、2的排序,重写当前数组。
你能想出一个仅使用常数空间的一趟扫描算法吗?
解题思路
- 扫描一遍,记录不同元素的个数,然后填充数组,这肯定可行,但是不是考察目的
- 3个指针,r,w,b初始化为0,0,n-1,分别指向red的右侧,white的右侧,blue的左侧
- 如果w当前指向red,则交换r和w的元素,然后二者都+=1
- 因为r在w的左侧,所以交换时,只要r<w,则r指向的只可能是white,所以交换后r+=1
- 如果w当前指向white,则w+=1
- 如果w当前指向blue,则交换b和w的元素,然后b-=1
- 因为b在w的右侧,所以交换后,w指向的仍然可能是3种可能,所有b-=1后,w先别动
- 循环直到w到了b的右边(while w <= b)
代码
class Solution:
def sortColors(self, nums: List[int]) -> None:
"""
Do not return anything, modify nums in-place instead.
"""
# - sanity check
if not nums:
return nums
# - 3 pointers
# - r: right side of red
# - w: right white of white
# - b: left side of blue
r,w,b = 0,0,len(nums)-1
while w <= b:
if nums[w] == 0:
if w > r:
nums[r], nums[w] = nums[w], nums[r]
w += 1
r += 1
elif nums[w] == 1:
w += 1
elif nums[w] == 2:
if w < b:
nums[b], nums[w] = nums[w], nums[b]
b -= 1