Sort Colors — Solution in Python, Java & C++
Sort Colors solution explained: step-by-step approach, complexity and code in Python, Java, C++ and JavaScript. A medium problem on arrays and sorting.
Sort Colors is a medium coding problem on arrays, sorting and two pointers, asked in Morgan Stanley and Mphasis interviews.
- Difficulty: Medium
- Topics: Arrays, Sorting, Two Pointers
- Asked at: Morgan Stanley, Mphasis
- Time limit: 2 s
- Memory limit: 256 MB
- Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby
Problem statement
Given an array nums with n objects colored red (0), white (1), or blue (2), sort them in place so that objects of the same color are adjacent, in the order red, white, blue, and return the array.
Solve it without the library sort — ideally in one pass with constant extra space (the Dutch National Flag algorithm).
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
1 <= nums.length <= 30nums[i] is 0, 1, or 2.
Hints for Sort Colors
Hint 1
Three regions: 0s on the left, 2s on the right, 1s in the middle.
Hint 2
Pointers low/mid/high — swap 0s to low, 2s to high, walk past 1s.
How to solve Sort Colors
With only three possible values you can sort in a single pass using the Dutch National Flag partition. Maintain three regions — settled 0s at the front, settled 2s at the back, and 1s in between — and grow them with one scanning pointer and two boundaries.
Approach
- Keep
low(end of the zeros region),mid(the scanner) andhigh(start of the twos region). - While
mid <= high, look atnums[mid]. - If it is
0, swap it into the zeros region: swap withlow, then advance bothlowandmid. - If it is
1, it already belongs in the middle — just advancemid. - If it is
2, swap it back tohighand decrementhigh, without advancingmid— the value swapped in is unexamined.
Why it works
The invariant is that everything before low is 0, everything between low and mid is 1, and everything after high is 2. Each branch restores it while shrinking the unexamined region [mid, high]. The asymmetry is the crux: a swap with low brings back a value already known to be 1 (or mid == low), so it is safe to step over, whereas a swap with high brings back something never inspected, which must be re-examined.
Complexity
- Time —
O(n) - Space —
O(1)
Pitfalls
- Do not advance
midafter swapping withhigh— the incoming value is unclassified and skipping it leaves the array unsorted. - The loop condition is
mid <= high, inclusive; stopping at<leaves the final element unplaced. - Counting the three values and rewriting the array also works and is arguably simpler, but takes two passes.
- In languages with unsigned indices,
highcan fall below zero — keep the indices signed.
Reference solution
Python
from typing import List
def sortColors(nums: List[int]) -> List[int]:
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] == 2:
nums[mid], nums[high] = nums[high], nums[mid]
high -= 1
else:
mid += 1
return numsJava
public static int[] sortColors(int[] nums) {
int low = 0, mid = 0, high = nums.length - 1;
while (mid <= high) {
if (nums[mid] == 0) {
int t = nums[low];
nums[low] = nums[mid];
nums[mid] = t;
low++;
mid++;
} else if (nums[mid] == 1) {
mid++;
} else {
int t = nums[high];
nums[high] = nums[mid];
nums[mid] = t;
high--;
}
}
return nums;
}C++
vector<int> sortColors(vector<int>& nums) {
int low = 0, mid = 0, high = (int) nums.size() - 1;
while (mid <= high) {
if (nums[mid] == 0) {
swap(nums[low], nums[mid]);
low++;
mid++;
} else if (nums[mid] == 1) {
mid++;
} else {
swap(nums[mid], nums[high]);
high--;
}
}
return nums;
}JavaScript
var sortColors = function(nums) {
let low = 0, mid = 0, high = nums.length - 1;
while (mid <= high) {
if (nums[mid] === 0) {
const t = nums[low]; nums[low] = nums[mid]; nums[mid] = t;
low++; mid++;
} else if (nums[mid] === 2) {
const t = nums[mid]; nums[mid] = nums[high]; nums[high] = t;
high--;
} else {
mid++;
}
}
return nums;
};Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.
All 1199 arrays problems · the whole catalogue
Learn the technique: Arrays · Sorting Algorithms