Sort Colors II

Given an array of n objects with k different colors (numbered from 1 to k), sort them so that objects of the same color are adjacent, with the colors in the order 1, 2, ... k.

http://www.lintcode.com/en/problem/sort-colors-ii/

Discussion

A rather straight forward solution is a two-pass algorithm using counting sort. That will cost O(k) extra memory.

Solution

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

还有个partition的实现方法，TBD。。。

调用partition，也可以自己实现

class Solution{
public:
    /**
     * @param colors: A list of integer
     * @param k: An integer
     * @return: nothing
     */    
    void sortColors2(vector<int> &colors, int k) {
        // write your code here
        int n = colors.size();
        if(n == 0) return;
        vector<int>::iterator start = colors.begin();
        for(int i=1; i<=k; i++) {
            start = partition(start, colors.end(), bind1st(equal_to<int>(),i));
        }
    }
};

原地计数排序

// Time:  O(n)
// Space: O(1)

// Inplace counting sort.
class Solution{
public:
    /**
     * @param colors: A list of integer
     * @param k: An integer
     * @return: nothing
     */
    void sortColors2(vector<int> &colors, int k) {
        for (int i = 0; i < colors.size(); ++i) {
            if (colors[i] > 0) {
                int pos = colors[i] - 1;
                if (colors[pos] <= 0) {  // Bucket exists.
                    --colors[pos];
                    colors[i] = 0;
                }
                else {  // Init a new bucket.
                    colors[i] = colors[pos];
                    colors[pos] = -1;
                    --i;
                }
            }
        }

        for (int i = colors.size() - 1, pos = k - 1; pos >= 0; --pos) {
            while (colors[pos] < 0) {  // Reorder the color by count of each bucket.
                ++colors[pos];
                colors[i--] = pos + 1;
            }
        }
    }
};