三色排序练习

有一个只由0，1，2三种元素构成的整数数组，请使用交换、原地排序而不是使用计数进行排序。
给定一个只含0，1，2的整数数组A及它的大小，请返回排序后的数组。保证数组大小小于等于500。
测试样例：
[0,1,1,0,2,2],6
返回：[0,0,1,1,2,2]

class ThreeColor {public:    vector<int> sortThreeColor(vector<int> A, int n) {        int zero_after=0;        int two_before=n-1;        int i=0;        while(i<=two_before)            {            if(A[i]==0)            {                swapnum(A[i],A[zero_after]);                ++zero_after;                ++i;            }            else if(A[i]==1)                {                ++i;            }            else if(A[i]==2)             {                  swapnum(A[i],A[two_before]);                  --two_before;             }        }        return A;    }    void swapnum(int &a,int &b)        {            int temp=a;            a=b;            b=temp;        }};