Actually we can have a code that works for arbitrarily many colors, as long as the colors are numbered with consecutive integers starting at 0. The number of colors does not even has to be part of the input. Try it out! Have fun.

```
public class Solution {
public void sortColors(int[] nums) {
int color = 0;
int i = 0;
while(i < nums.length) {
if(nums[i] != color) {
int j = i + 1;
while(j < nums.length && nums[j] != color) {
j++;
}
if(j < nums.length) {
nums[j] = nums[i];
nums[i] = color;
i++;
}else {
color++;
}
}else {
i++;
}
}
}
}
```