75. sort colors


Leetcode
Two Pointers
Array
Sort

Given an array with $n$ objects colored red, white or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white and blue.

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.

Note: You are not suppose to use the library's sort function for this problem.

Example:

Input: [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]

Follow up:

• A rather straight forward solution is a two-pass algorithm using counting sort.
• First, iterate the array counting number of 0's, 1's, and 2's, then overwrite array with total number of 0's, then 1's and followed by 2's.
• Could you come up with a one-pass algorithm using only constant space?

分析¶

题目都说了可以使用计数排序(couting sort)，那么肯定先用计数排序：先计数每一种颜色的数目，然后按顺序将数组填充上相应数目的颜色。

public void sortColors(int[] nums) {
    int[] colorCount = new int[3]; // count 0, 1, 2
    for (int color : nums)
        colorCount[color]++;
    for (int i = 0, index = 0; i < 3; i++)
        for ( ; colorCount[i] > 0; colorCount[i]--)
            nums[index++] = i;
}

这道题目实际上是Dutch national flag problem，是鼎鼎大名的Dijkstra提出的。方法是使用3-Way Partition。平均时间复杂度是$O(n)$，空间复杂度是$O(1)$.

// 3-way partition
public void sortColors(int[] nums) {
    int lt = 0, i = 0, gt = nums.length - 1;
    int v = 1;
    while (i <= gt) {
        if (nums[i] < v) swap(nums, lt++, i++);
        else if (nums[i] > v) swap(nums, i, gt--);
        else i++;
    }
}
private void swap(int[] nums, int i, int j) {
    int tmp = nums[i];
    nums[i] = nums[j];
    nums[j] = tmp;
}