给定一个包含红色、白色和蓝色,一共 n 个元素的数组,原地对它们进行排序,使得相同颜色的元素相邻,并按照红色、白色、蓝色顺序排列。
此题中,我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。
注意:
不能使用代码库中的排序函数来解决这道题。
示例:
输入: [2,0,2,1,1,0]
输出: [0,0,1,1,2,2]
进阶:
- 一个直观的解决方案是使用计数排序的两趟扫描算法。
首先,迭代计算出0、1 和 2 元素的个数,然后按照0、1、2的排序,重写当前数组。 - 你能想出一个仅使用常数空间的一趟扫描算法吗?
解题思路
首先想到的思路就是写一个快速排序(因为不能使用系统自带的排序函数),所以我就写了一个三路快排(没有使用random
)
class Solution:
def _sortColors(self, nums, l, r):
if r <= l: return
lt = l
gt = r + 1
i = l + 1
while i < gt:
if nums[i] < nums[l]:
lt += 1
nums[i], nums[lt] = nums[lt], nums[i]
i += 1
elif nums[i] > nums[l]:
gt -= 1
nums[i], nums[gt] = nums[gt], nums[i]
else:
i += 1
nums[l], nums[lt] = nums[lt], nums[l]
self._sortColors(nums, l, lt - 1)
self._sortColors(nums, gt, r)
def sortColors(self, nums):
"""
:type nums: List[int]
:rtype: void Do not return anything, modify nums in-place instead.
"""
self._sortColors(nums, 0, len(nums) - 1)
我们知道这个算法的时间复杂度是O(nlogn)
,那么我们能不能使用O(n)
就解决这个问题呢?我们一直都忽略了题目的条件,我们没有充分利用它(只有0,1,2
三个元素)
我们可以这么想,我们先统计总共有多少个0
、1
和2
,然后再将这些值放回就可以了,这实际上是计数排序的原理。
class Solution:
def sortColors(self, nums):
"""
:type nums: List[int]
:rtype: void Do not return anything, modify nums in-place instead.
"""
count = [0, 0, 0]
for i in nums:
assert(i >= 0 and i <= 2)
count[i] += 1
index = 0
for i, num in enumerate(count):
for j in range(num):
nums[index] = i
index += 1
我们通过分析发现这个算法的时间复杂度是O(n)
级别的,而空间复杂度和我们统计的数的个数有关O(3)
。并且我们这个解法中扫描了两遍数组,我们有没有只扫描一遍的做法呢?
我们是否可以使用三路快排的思路(保证遍历数组的过程中0
位于数组的开头,1
位于数组的中间,2
位于数组的末尾)
class Solution:
def sortColors(self, nums):
"""
:type nums: List[int]
:rtype: void Do not return anything, modify nums in-place instead.
"""
zero = -1 # [0, zero] 0
i = 0 # [zero + 1, i) 1
two = len(nums) # [two, len(nums) - 1] 2
while i < two:
if nums[i] == 1:
i += 1
elif nums[i] == 2:
two -= 1
nums[two], nums[i] = nums[i], nums[two]
else:
# assert(nums[i] == 0)
zero += 1
nums[zero], nums[i] = nums[i], nums[zero]
i += 1
该问题的其他语言版本添加到了我的GitHub Leetcode
如有问题,希望大家指出!!!