LeetCode - Sort Colors
Problem statement
Given an array nums with n objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue.
We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.
You must solve this problem without using the library's sort function.
Problem statement taken from: https://leetcode.com/problems/sort-colors
Example 1:
Input: nums = [2, 0, 2, 1, 1, 0]
Output: [0, 0, 1, 1, 2, 2]
Example 2:
Input: nums = [2, 0, 1]
Output: [0, 1, 2]
Example 3:
Input: nums = [0]
Output: [0]
Example 4:
Input: nums = [1]
Output: [1]
Constraints:
- n == nums.length
- 1 <= n <= 300
- nums[i] is 0, 1, or 2
Explanation
Simple counting
The simple approach will be to count the occurrences of each integer 0, 1 and 2 using three different variables.
Using the above three count variables we fill in the array.
A small C++ snippet of the approach will look like this:
for (int i = 0; i < n; i++) {
if (arr[i] == 0)
count0++;
if (arr[i] == 1)
count1++;
if (arr[i] == 2)
count2++;
}
for (int i = 0; i < count0; i++)
arr[i] = 0;
for (int i = count0; i < (count0 + count1); i++)
arr[i] = 1;
for (int i = (count0 + count1); i < n; i++)
arr[i] = 2;The time complexity of the above program is O(N). But in the above approach, we iterate the array twice.
Dutch national flag problem
We can use the approach of Dutch national flag problem. The problem was posed with three colors, here we have 0, 1, and 2. The array is divided into four sections to solve this problem.
Let's check the algorithm:
- Keep three indices low = 1, mid = 1 and, high = N and, there are four ranges,
1 to low (the range containing 0),
low to mid (the range containing 1),
mid to high (the range containing unknown elements) and
high to N (the range containing 2).
- Traverse the array from start to end and mid is less than high. (Loop counter is i)
- if element == 0
- swap the element with the element at index low and update low = low + 1 and mid = mid + 1
- if element == 1
- set mid = mid + 1
- if element == 2
- swap the element with the element at index high and update high = high – 1.
- set i = i – 1.
- return array
The time complexity of the program is O(N) as we iterate the array only once. Space complexity is O(1) because we do not use any other additional data structures.
C++ solution
class Solution {
public:
void sortColors(vector<int>& nums) {
int low = 0, mid = 0, high = nums.size() - 1;
while(mid <= high){
switch (nums[mid]){
case 0:
swap(nums[low++], nums[mid++]);
break;
case 1:
mid++;
break;
case 2:
swap(nums[mid], nums[high--]);
break;
}
}
}
};Golang solution
func sortColors(nums []int) {
low := 0
mid := 0
high := len(nums) - 1
for mid <= high {
switch (nums[mid]) {
case 0:
tmp := nums[low]
nums[low] = nums[mid]
nums[mid] = tmp
low++
mid++
break
case 1:
mid++
break
case 2:
tmp := nums[mid]
nums[mid] = nums[high]
nums[high] = tmp
high--
break
}
}
}Javascript solution
var sortColors = function(nums) {
function swap(i, j) {
[nums[i], nums[j]] = [nums[j], nums[i]];
}
let low = 0;
let high = nums.length - 1;
let mid = 0;
while (mid <= high) {
const n = nums[mid];
if (n === 0) {
swap(mid, low);
low++;
mid++;
} else if (n === 2) {
swap(mid, high);
high--;
} else {
mid++;
}
}
};Let's dry run the problem
Input: nums = [2, 0, 2, 1, 1, 0]
Step 1: low = 0
mid = 0
high = nums.length() - 1
= 6 - 1
= 5
Step 2: loop while mid < = high
0 <= 5
true
switch (nums[mid])
nums[mid] = nums[0]
= 2
case 2:
swap(nums[mid], nums[high--])
swap(nums[0], nums[5])
swap(2, 0)
nums = [0, 0, 2, 1, 1, 2]
high--
high = 4
Step 3: loop while mid < = high
0 <= 4
true
switch (nums[mid])
nums[mid] = nums[0]
= 0
case 0:
swap(nums[low++], nums[mid++])
swap(nums[0], nums[0])
swap(0, 0)
nums = [0, 0, 2, 1, 1, 2]
low++
mid++
low = 1
mid = 1
Step 4: loop while mid < = high
1 <= 4
true
switch (nums[mid])
nums[mid] = nums[1]
= 0
case 0:
swap(nums[low++], nums[mid++])
swap(nums[1], nums[1])
swap(1, 1)
nums = [0, 0, 2, 1, 1, 2]
low++
mid++
low = 2
mid = 2
Step 5: loop while mid < = high
2 <= 4
true
switch (nums[mid])
nums[mid] = nums[2]
= 2
case 2:
swap(nums[mid], nums[high--])
swap(nums[2], nums[4])
swap(2, 1)
nums = [0, 0, 1, 1, 2, 2]
high--
high = 3
Step 6: loop while mid < = high
2 <= 3
true
switch (nums[mid])
nums[mid] = nums[2]
= 1
case 1:
mid++
mid = 3
Step 7: loop while mid < = high
3 <= 3
true
switch (nums[mid])
nums[mid] = nums[3]
= 1
case 1:
mid++
mid = 4
Step 8: loop while mid < = high
4 <= 3
false
The result is [0, 0, 1, 1, 2, 2]