2_17三色排序

有一个只由0，1，2三种元素构成的整数数组，请使用交换、原地排序而不是使用计数进行排序。

给定一个只含0，1，2的整数数组A及它的大小，请返回排序后的数组。保证数组大小小于等于500。

测试样例：
输入：[0,1,1,0,2,2],6
返回：[0,0,1,1,2,2]

class ThreeColor {
public:
    void swap(vector<int> &A, int i, int j)
    {
        int temp = A[i];
        A[i] = A[j];
        A[j] = temp;
    }
    const static int first = 0, mid = 1, last = 2;
    vector<int> sortThreeColor(vector<int> A, int n) {
        // write code here
        int cmp_val = 1;
        int p_first = 0;
        int p_last = n-1;
        int p_curr = 0;
        for(int i=0; i<n; i++){
            switch(A[p_curr])
            {
                case first:
                    swap(A, p_curr++, p_first++);
                    break;
                case mid:
                    ++p_curr;
                    break;
                case last:
                    swap(A, p_curr, p_last--);
                    break;
                default:
                    break;
            }
        }
        return A;
    }
};