LeetCode--Remove Duplicates from Sorted Array II

Follow up for "Remove Duplicates":
What if duplicates are allowed at most twice?

For example,
Given sorted array A = [1,1,1,2,2,3],

Your function should return length = 5, and A is now [1,1,2,2,3].

class Solution 
{
public:
    int removeDuplicates(int A[], int n) 
    {
        if(n==0 || n==1 || n==2)
            return n;
        int loc=1;
        int count = 1;
        int flag = A[0];
        for(int i=1; i<n&&A[i]>=A[i-1];i++)
        {
            if(A[i] == flag)
            {
                count++;
                if(count < 3)
                {
                    swap(A,i,loc);
                    loc++;
                }
            }
            else
            {
                flag = A[i];
                swap(A,i,loc);
                count = 1;
                loc++;
            }
        }
        return loc;
        
    }
    void swap(int* a, int i, int j)
    {
        int t = a[i];
        a[i] = a[j];
        a[j] = t;
    }
};

LeetCode--Remove Duplicates from Sorted Array II

原文：http://blog.csdn.net/shaya118/article/details/42673723