【Lintcode】143. Sort Colors II

题目地址：

https://www.lintcode.com/problem/sort-colors-ii/description

给定一个数组$A$，再给定一个数$k$，$A$里面的数字都是$1\sim k$。要求将$A$排序。

法1：计数排序。开一个长度为$k$的数组，记录一下$1\sim k$每个数出现的次数，然后再把$A$填起来即可。代码如下：

public class Solution {
    /**
     * @param colors: A list of integer
     * @param k:      An integer
     * @return: nothing
     */
    public void sortColors2(int[] colors, int k) {
        // write your code here
        if (colors == null || colors.length == 0) {
            return;
        }
        
        int[] count = new int[k];
        for (int color : colors) {
            count[color - 1]++;
        }
    
        int idx = 0;
        for (int i = 0; i < count.length; i++) {
            for (int j = 0; j < count[i]; j++) {
                colors[idx++] = i + 1;
            }
        }
    }
}

时间复杂度$O(n)$，空间$O(k)$。

法2：彩虹排序。类似于快速排序里的partition做法，取一个pivot，然后把小于和大于pivot的数放到两边，这样每次partition完都能确定一个数已经全部放到中间了（也就是放到分界点上了）。彩虹排序比普通快速排序的写法更优的点在于，当同一个数字出现多次的时候，彩虹排序排完一遍之后会把这个数字都连续的存放到中间。所以在重复数字非常多的时候，彩虹排序会更快一些。代码如下：

public class Solution {
    /**
     * @param colors: A list of integer
     * @param k: An integer
     * @return: nothing
     */
    public void sortColors2(int[] colors, int k) {
        // write your code here
        partition(colors, 0, colors.length - 1);
    }
    
    private void partition(int[] colors, int l, int r) {
        if (l >= r) {
            return;
        }
        
        int midCol = colors[l + (r - l >> 1)];
        // 开始彩虹排序，将小于midCol的数都放到左边，大于mideCol的数都放到右边；
        // 维护两个指针i和k都指向小于和大于midCol的数中最中间的那个的下标
        int i = l - 1, j = l, k = r + 1;
        while (j < k) {
            if (colors[j] < midCol) {
                swap(colors, ++i, j++);
            } else if (colors[j] > midCol) {
                swap(colors, --k, j);
            } else {
                j++;
            }
        }
        
        // 递归排序左右两边
        partition(colors, l, i);
        partition(colors, k, r);
    }
    
    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

时间复杂度$O(n\log k)$，空间$O(\log k)$（递归栈深度。这里指的都是平均复杂度）。