Sort Colors
An array nums holds a row of objects, each painted one of three colours, encoded as 0 (red), 1 (white) and 2 (blue).
Rearrange the array in place so that all reds come first, then all whites, then all blues; in other words, sort it in non-decreasing order. Your function should not return anything: the judge inspects nums after your function finishes.
Do not call a built-in sort. Try to do it in a single pass using only constant extra memory.
Example 1
Example 2
Example 3
Constraints
1 <= nums.length <= 300- nums[i] is 0, 1 or 2
Pattern Path · Two Pointers · problem 20 of 169
low (next slot for a 0), mid (current element) and high (next slot for a 2). After swapping a 2 to the back, do not advance mid: the value swapped in has not been examined yet.Start simple
Count the 0s, 1s and 2s, then rewrite the array: O(n) time and O(1) space, but two passes (or a general sort at O(n log n)).
The key idea
Maintain four regions: nums[0..low) are all 0, nums[low..mid) are all 1, nums[mid..high] is still unexamined, and nums(high..end] are all 2. Initially everything is unexamined.
Look at nums[mid]:
- 0: swap it with
nums[low]. The value that comes back is a 1 (or the same 0 when the 1-zone is empty), so bothlowandmidmove forward. - 1: it is already in the right zone; advance
mid. - 2: swap it with
nums[high]and shrinkhigh. The element swapped intomidcame from the unexamined region, somidstays put and it is checked next.
Each step shrinks the unexamined region by one, so the loop ends after at most n steps with every element in its zone. This is the classic “Dutch national flag” partition.
Complexity
- Time:
O(n) - Space:
O(1)
Python
class Solution:
def sortColors(self, nums):
# nums[:low] are 0s, nums[low:mid] are 1s, nums[high+1:] are 2s.
low, mid, high = 0, 0, len(nums) - 1
while mid <= high:
if nums[mid] == 0:
nums[low], nums[mid] = nums[mid], nums[low]
low += 1
mid += 1
elif nums[mid] == 1:
mid += 1
else:
nums[mid], nums[high] = nums[high], nums[mid]
high -= 1 # the swapped-in value is unexamined, so mid stays
JavaScript
var sortColors = function(nums) {
// nums[0..low) are 0s, nums[low..mid) are 1s, nums(high..] are 2s.
let low = 0, mid = 0, high = nums.length - 1;
while (mid <= high) {
if (nums[mid] === 0) {
[nums[low], nums[mid]] = [nums[mid], nums[low]];
low++;
mid++;
} else if (nums[mid] === 1) {
mid++;
} else {
[nums[mid], nums[high]] = [nums[high], nums[mid]];
high--; // the swapped-in value is unexamined, so mid stays
}
}
};
What the interviewer is checking
This checks whether you can design and reason about loop invariants for multiple pointers, and explain exactly when each pointer may move.
Mistakes to avoid
- Advancing
midafter swapping withhigh, which leaves an unchecked 0 or 2 in the middle. - Looping while
mid < highinstead ofmid <= high, which skips the last unexamined element. - Returning a new array instead of modifying
numsin place.
Follow-ups to expect
- What if there are k colours instead of three? (Counting sort, or repeated partitioning.)
- Can you make the partition stable for records carrying extra data? (Not in place; use buckets.)
- How does this relate to quicksort with many duplicates? (Three-way partitioning around the pivot.)
In real software
Three-way partitioning is how robust quicksort implementations handle many equal keys, and it is used to bucket items by priority level (high, normal, low) without extra memory.