Sort colors, and sort k colors, C++ solution


  • 3
    X

    Solution:

    //3 color sort
    //2 partition indexes
    //time complexity: O(N)
    //space complexity: O(1)
    class Solution{
    public:
        /**
         * @param nums: A list of integer which is 0, 1 or 2 
         * @return: nothing
         */    
        void sortColors(vector<int> &nums) {
            // write your code here
            if (nums.size() == 0) return;
            int zeroPartition = 0, twoPartition = nums.size() - 1;
            for (int i = 0; i <= twoPartition; ) {
                if (nums[i] == 0) {
                    swap(nums[zeroPartition], nums[i]); zeroPartition++; i++;
                } else if (nums[i] == 2) {
                    swap(nums[twoPartition], nums[i]); twoPartition--;
                } else {
                    i++;
                }
            }
        }
        
        void swap(int& a, int& b) {
            int tmp = a; a = b; b = tmp;
        }
    };
    

    There is also an extension question on lintcode: sort k colors. http://www.lintcode.com/en/problem/sort-colors-ii/#

    Basically, the naive solution is a two-pass algorithm using counting sort. That will cost O(k) extra memory. How can you solve it without using extra memory?

    The O(1) space solution is just an extension of sort colors:

    //sort k colors, extend by sort 3 colors
    class Solution{
    public:
        /**
         * @param colors: A list of integer
         * @param k: An integer
         * @return: nothing
         */    
        void sortColors2(vector<int> &colors, int k) {
            if (colors.size() == 0) return;
            // write your code here
            int lowColor = 1, highColor = k;
            int lpartition = 0, hpartition = colors.size() - 1;
            while (lowColor < highColor) {
                for (int i = lpartition; i <= hpartition; ) {
                    if (colors[i] == lowColor) {
                        swap(colors[i], colors[lpartition]); lpartition++; i++;
                    } else if (colors[i] == highColor) {
                        swap(colors[i], colors[hpartition]); hpartition--;
                    } else {
                        i++;
                    }
                }
                lowColor++; highColor--;
            }
        }
    };
    

    Can someone talks about this solution's time complexity? Thanks


  • 1

    It will be O(kn/2) since the inner loop will be O(N) while the outer O(k/2) -> linear complexity, but when it comes to the details of it, I must say it will be less than this since the start position will be higher and higher and the ending position will be lower and lower, then in average it will be about O(kn/4). By the way, nice solution!


  • 0

    Great, really learn a lot. Thanks.


  • 0
    P

    Since i is going get set back to lpartition inside the while loop, in the worst case inner (for) loop is going to run 'n' times (i.e. when n==k). I would say worst case complexity would be O(n^2), or simply O(n*k)


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.