LeetCode 75 - Sort Colors

作者 QIFAN 日期 2017-01-19
LeetCode 75 - Sort Colors

原题链接: 75. Sort Colors


题干

给定一个数组包含 0 ,1 ,2 ,分别代表红色、白色、蓝色,要求对着数组排序,使相同颜色相邻,且顺序为红、白、蓝。

思路

第一想法就是直接用 sort 方法,时间复杂度为 $O(nlogn) $ 。

但是呢这里只有三种颜色,实际上已经是部分排序了。我想到了快速排序中的 partition 。partition 后将数组分为小于 pivot 与大于 pivot 的部分。所以应用带该题,pivot = 1 。然后怎么让 1 集中到中间?我想到用三个指针,分别代表 lo ,hi ,mid 。

  • lo 从头开始,向后移动 lo 直到找到大于等于 pivot
  • hi 从尾开始向前移动 hi 直到找到小于等于 pivot
  • mid 从三分之一处开始,向后移动直到不等于 pivot
  • 相互交换到正确位置
  • 重复上述步骤, lo 到达三分之一处时停止,hi 和 mid 相交时停止。

时间复杂度 $O(n)$ ,空间复杂度 $O(1) $ 。这个方法判断条件太多了,所以直接使用两次 partition 。

  • 整个数组以 0 为 pivot 进行 partition ,将所有 0 归到最前,并返回 0 的末尾位置 zeroEnd 。
  • zeroEnd 到数组末尾此时只包含了 1 和 2 ,再将此部分以 1 为 pivot 进行 partition 。

时间复杂度为 $O(n)$ 。

坑的地方在于:

  1. 我 partition 返回的是 0 的末尾,要 + 1 才是另一部分的开头
  2. 如果数组中不包含 0 ,那同样返回 0 ,此时判断一下如果第一位是 0 再加 1 。

代码:

public void sortColors(int[] nums) {
if (nums.length <= 1) {
return;
}
// 第一次 partition
int zeroEnd = partition(nums, 0, nums.length - 1, 0);
if (zeroEnd == nums.length) {
return;
}
if (nums[zeroEnd] == 0) {
zeroEnd++;
}
// 第二次 partition
partition(nums, zeroEnd, nums.length - 1, 1);
}
private int partition(int[] nums, int lo, int hi, int pivot) {
while (lo < hi) {
while (lo < nums.length && nums[lo] == pivot) {
lo++;
}
while (hi >= 0 && nums[hi] > pivot) {
hi--;
}
if (lo <= hi) {
swap(nums, lo++, hi--);
}
}
return lo;
}
private void swap(int[] nums, int a, int b) {
int temp = nums[a];
nums[a] = nums[b];
nums[b] = temp;
}