Sort Colors
MedPrompt
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.
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]
Constraints:
n == nums.length1 <= n <= 300nums[i]is either0,1, or2.
Follow up: Could you come up with a one-pass algorithm using only constant extra space?
Approaches
3 approaches with complexity analysis and trade-offs.
The most straightforward, albeit inefficient, way to solve this problem is to use a standard comparison-based sorting algorithm. Since the problem prohibits using the library's built-in sort function, we can implement a simple one like Bubble Sort. This approach repeatedly steps through the list, compares adjacent elements, and swaps them if they are in the wrong order. The passes through the list are repeated until the list is sorted.
Algorithm
- Implement a standard sorting algorithm like Bubble Sort.
- Use nested loops to iterate through the array.
- The outer loop runs from
i = 0ton-2. - The inner loop runs from
j = 0ton-i-2. - Inside the inner loop, compare
nums[j]andnums[j+1]. - If
nums[j] > nums[j+1], swap the two elements. - After the loops complete, the array will be sorted.
Walkthrough
This approach treats the problem as a generic sorting problem. We can implement an algorithm like Bubble Sort, which works by repeatedly swapping adjacent elements if they are in the wrong order. We would need two nested loops. The outer loop decreases the effective size of the array to be sorted in each iteration, as the largest elements 'bubble up' to the end. The inner loop performs the actual comparisons and swaps.
For an input like [2,0,2,1,1,0], the algorithm would perform multiple passes, gradually moving the 0s to the left and 2s to the right, until the array is fully sorted as [0,0,1,1,2,2].
class Solution { public void sortColors(int[] nums) { int n = nums.length; boolean swapped; for (int i = 0; i < n - 1; i++) { swapped = false; for (int j = 0; j < n - i - 1; j++) { if (nums[j] > nums[j + 1]) { // Swap nums[j] and nums[j+1] int temp = nums[j]; nums[j] = nums[j + 1]; nums[j + 1] = temp; swapped = true; } } // If no two elements were swapped by inner loop, then break if (!swapped) { break; } } }}Complexity
Time
O(n^2)
Space
O(1)
Trade-offs
Pros
Easy to understand if familiar with basic sorting algorithms.
It's an in-place sorting algorithm, so it uses constant extra space.
Cons
Very inefficient compared to other methods.
Does not leverage the specific properties of the input (only three distinct values).
Has a time complexity of O(n^2), which is slow for larger inputs.
A more efficient approach is to use a variation of Counting Sort. Since we know the range of values is limited to 0, 1, and 2, we can count the occurrences of each color in a first pass. Then, in a second pass, we can overwrite the original array with the correct number of 0s, followed by 1s, and then 2s. This is a two-pass algorithm.
Algorithm
- Initialize three counters:
count0 = 0,count1 = 0,count2 = 0. - First Pass: Iterate through the
numsarray from start to end. - For each element, increment the corresponding counter. If
num == 0, incrementcount0, etc. - Second Pass: Overwrite the
numsarray. - Initialize a pointer
i = 0. - Fill the first
count0elements ofnumswith0. - Fill the next
count1elements ofnumswith1. - Fill the final
count2elements ofnumswith2.
Walkthrough
This method involves two distinct steps. First, we iterate through the entire array to count how many 0s, 1s, and 2s are present. We can use three integer variables, say count0, count1, and count2, for this purpose.
After counting, we know the exact composition of the sorted array. For example, if count0=2, count1=2, and count2=2, the sorted array must be [0,0,1,1,2,2]. The second step is to modify the input array in-place. We use a pointer, starting at index 0, and fill the array first with count0 zeros, then with count1 ones, and finally with count2 twos.
class Solution { public void sortColors(int[] nums) { int count0 = 0, count1 = 0, count2 = 0; for (int num : nums) { if (num == 0) { count0++; } else if (num == 1) { count1++; } else { count2++; } } int i = 0; // Overwrite array with count0 zeros for (int j = 0; j < count0; j++) { nums[i++] = 0; } // Overwrite array with count1 ones for (int j = 0; j < count1; j++) { nums[i++] = 1; } // Overwrite array with count2 twos for (int j = 0; j < count2; j++) { nums[i++] = 2; } }}Complexity
Time
O(n)
Space
O(1)
Trade-offs
Pros
Linear time complexity O(n), which is a significant improvement over O(n^2).
Simple logic to implement and understand.
Uses constant extra space (O(1)) for the counters.
Cons
Requires two passes over the array, which is less optimal than a single-pass solution.
It's not a true in-place sort as it overwrites the array based on counts rather than swapping elements.
The optimal solution is a one-pass, in-place algorithm often referred to as the Dutch National Flag problem. This approach uses three pointers (low, mid, and high) to partition the array into three sections: a section for 0s at the beginning, a section for 2s at the end, and a section for 1s in the middle. We iterate through the array with the mid pointer, swapping elements into their correct partitions.
Algorithm
- Initialize three pointers:
low = 0,mid = 0, andhigh = n - 1. - The region
0tolow-1contains0s. - The region
lowtomid-1contains1s. - The region
high+1ton-1contains2s. - The region
midtohighis the unsorted part. - Iterate while
mid <= high:- If
nums[mid] == 0: Swapnums[low]withnums[mid]. Increment bothlowandmid. - If
nums[mid] == 1: The element is in its correct partition. Incrementmid. - If
nums[mid] == 2: Swapnums[mid]withnums[high]. Decrementhigh. Do not incrementmidas the newnums[mid]needs to be processed.
- If
Walkthrough
This elegant one-pass solution partitions the array in-place. We maintain three pointers:
low: Marks the boundary of the0's section. All elements beforeloware0.high: Marks the boundary of the2's section. All elements afterhighare2.mid: Is the current element being processed.
We iterate with mid from the beginning of the array.
- If
nums[mid]is0, it belongs to thelowsection. We swap it withnums[low]and increment bothlowandmid. - If
nums[mid]is1, it's in the correct place for now, so we just move to the next element by incrementingmid. - If
nums[mid]is2, it belongs to thehighsection. We swap it withnums[high]and decrementhigh. We do not incrementmidin this case, because the element we just swapped into themidposition from thehighend is unprocessed and needs to be checked.
The loop continues until mid crosses high, at which point the array is fully sorted.
class Solution { public void sortColors(int[] nums) { int low = 0; int mid = 0; int high = nums.length - 1; while (mid <= high) { switch (nums[mid]) { case 0: { int temp = nums[low]; nums[low] = nums[mid]; nums[mid] = temp; low++; mid++; break; } case 1: { mid++; break; } case 2: { int temp = nums[mid]; nums[mid] = nums[high]; nums[high] = temp; high--; break; } } } }}Complexity
Time
O(n)
Space
O(1)
Trade-offs
Pros
Optimal time complexity of O(n).
Optimal space complexity of O(1).
Achieves the sort in a single pass over the data.
Directly addresses the follow-up question.
Cons
The logic, especially the pointer movements, can be slightly tricky to get right on the first try.
Solutions
Solution
public class Solution { public void SortColors(int[] nums) { int i = -1, j = nums.Length, k = 0; while (k < j) { if (nums[k] == 0) { swap(nums, ++i, k++); } else if (nums[k] == 2) { swap(nums, --j, k); } else { ++k; } } } private void swap(int[] nums, int i, int j) { int t = nums[i]; nums[i] = nums[j]; nums[j] = t; }}class Solution {public void sortColors(int[] nums) { int i = -1, j = nums.length, k = 0; while (k < j) { if (nums[k] == 0) { swap(nums, ++i, k++); } else if (nums[k] == 2) { swap(nums, --j, k); } else { ++k; } } }private void swap(int[] nums, int i, int j) { int t = nums[i]; nums[i] = nums[j]; nums[j] = t; }}class Solution {public: void sortColors(vector<int> &nums) { int i = -1, j = nums.size(), k = 0; while (k < j) { if (nums[k] == 0) { swap(nums[++i], nums[k++]); } else if (nums[k] == 2) { swap(nums[--j], nums[k]); } else { ++k; } } }};class Solution : def sortColors ( self , nums : List [ int ]) -> None : i , j , k = - 1 , len ( nums ), 0 while k < j : if nums [ k ] == 0 : i += 1 nums [ i ], nums [ k ] = nums [ k ], nums [ i ] k += 1 elif nums [ k ] == 2 : j -= 1 nums [ j ], nums [ k ] = nums [ k ], nums [ j ] else : k += 1Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.