This is a standard problem in Computer Science and is known as Dutch National Flag Problem or Three way partition problem.
The idea is to maintain four sections with the following invariant:
- Section 1: nums[0…low-1] having only red color (denoted by 0)
- Section 2: nums[low…mid-1] having only white color (denoted by 1)
- Section 3: nums[mid…high] having unknown elements
- Section 4: nums[high+1…size(nums)-1] having only blue color (denoted by 2)
We keep three integer low, mid and high where low denotes the next available index for a 0 and high denotes the next available index for a 2. Go through the solution for a clearer idea.
Solution:
void sortColors(vector<int>& nums) {
int low = 0, mid = 0, high = nums.size()-1;
while (mid <= high) {
if (nums[mid] == 1)
mid++;
else if (nums[mid] == 0)
swap(nums[low++], nums[mid++]);
else
swap(nums[mid], nums[high--]);
}
}
Time complexity: O(n)
Space complexity: O(1)
Read more about two way partitioning and three way partitioning here: http://users.monash.edu/~lloyd/tildeAlgDS/Sort/Flag/