Given an array containing only
0,1, and2, sort it in-place so that all0scome first, followed by1s, and then2s.
Example:
Input: [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]
My First Thought (Counting Approach)
Since the array contains only three distinct values, I can count the occurrences of 0, 1, and 2 in one pass.
Then, overwrite the array by placing:
- All
0s - Then all
1s - Then all
2s
Code
class Solution {
public void sortColors(int[] nums) {
int zero = 0, one = 0, two = 0;
for (int num : nums) {
if (num == 0) zero++;
else if (num == 1) one++;
else two++;
}
int index = 0;
while (zero-- > 0) nums[index++] = 0;
while (one-- > 0) nums[index++] = 1;
while (two-- > 0) nums[index++] = 2;
}
}
Complexity
- Time Complexity: O(N)
- Space Complexity: O(1)
This works perfectly, but the follow-up asks for a one-pass solution.
The Key Observation
Instead of counting first and rewriting later, can we place elements directly into their correct regions while traversing the array?
Since there are only three values, we can maintain three sections:
0s | 1s | Unknown | 2s
This is where the famous Dutch National Flag Algorithm comes in.
Dutch National Flag Algorithm
We use three pointers:
low -> next position for 0
mid -> current element
high -> next position for 2
Rules
If nums[mid] == 0
Place it in the 0-region.
swap(low, mid);
low++;
mid++;
If nums[mid] == 1
Already in the correct region.
mid++;
If nums[mid] == 2
Place it in the 2-region.
swap(mid, high);
high--;
Notice that mid does not move here because the element swapped from the right side still needs to be processed.
Optimal Solution
class Solution {
public void sortColors(int[] nums) {
int low = 0;
int mid = 0;
int high = nums.length - 1;
while (mid <= high) {
if (nums[mid] == 0) {
swap(nums, low, mid);
low++;
mid++;
} else if (nums[mid] == 1) {
mid++;
} else {
swap(nums, mid, high);
high--;
}
}
}
private void swap(int[] nums, int i, int j) {
int temp = nums[i];
nums[i] = nums[j];
nums[j] = temp;
}
}
Complexity
- Time Complexity: O(N)
- Space Complexity: O(1)
What I Learned
The biggest takeaway wasn't the solution itself, but the pattern.
Whenever I see:
- Only 2 or 3 unique values
- In-place segregation
- One-pass requirement
I should immediately think about partitioning techniques and the Dutch National Flag Algorithm.
This pattern appears frequently in coding interviews and is much more useful than it initially seems.
I'm currently solving the Striver SDE Sheet Challenge and documenting my learnings, patterns, mistakes, and interview insights along the way.
🔗 LinkedIn: https://www.linkedin.com/in/jaspreetsodhi482/
🐦 X (Twitter): https://x.com/Jaspree54799902
💻 GitHub: https://github.com/codewithjaspreet
Top comments (0)