This problem is the same as an improved partition algorithm which puts the numbers equal to the pivot at the middle.

So basically, we can keep track the small part and equal part. Then scan once to partition the array.

```
void sortColors(vector<int>& nums) {
int small = 0;
int cntOne = 0;
for (int& num: nums) {
if (num == 1) {
swap(num, nums[small + cntOne]);
++cntOne;
} else if (num < 1) {
swap(num, nums[small + cntOne]);
swap(nums[small], nums[small + cntOne]);
++small;
}
}
}
```