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

• 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--;
}
}
};
``````

• 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)`