Leetcode: Sort Colors

// The easy way is to use count sort, but it is going to need two-pass
// The solution below need one-pass.
// The idea is to keep all the numbers before index0 are 0, and all the numbers after index2 are 2.
// One thing important is that, we should swap elements, instead of updating the array directly.
class Solution {
private:
void swap(int &a, int &b){
int tmp=a;
a=b;
b=tmp;
}
public:
void sortColors(int A[], int n) {
int index0=0, index2=n-1;
int current=0;
while(current<=index2){
if(A[current]==0) swap(A[index0++],A[current]);
else if(A[current]==2) swap(A[index2--],A[current]);
else current++;
current=max(current,index0);
}
}
};