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].

---

We can use the same idea as Remove Duplicates from Sorted Array. We can set up two indices, one for the old array, and one for the new array. If the current element is equal to the old one and the count of the element doesn't exceed 2, then we can move two indices one step forward. If the current element is equal to the old one and the count of the element exceeds 2, then just move the index for old array one step forward.

C++ Code:

/*
 * func: remove_duplicates
 * goal: remove duplicates from sorted array where duplicates are allowed at 
         most twice
 * @param A: input array
 * @param n: size of the array
 * return: the new length of the array
 */
/*
 * keep a new index of the new array while traversing
 * complexity: time O(n), space O(1)
 */
int remove_duplicates(int A[], int n){
    if(n <= 2){
        return n;
    }
    
    int size = 1;
    int count = 1;
    int old_idx = 1;
    int new_idx = 0;
    
    while(old_idx < n){
        //If current element is equal to previous one
        if(A[old_idx] == A[new_idx]){
            ++count;
            if(count <= 2){
                //If count is less than 2, move new_idx one step forward
                //new size plus one
                ++new_idx;
                ++size;
                //Copy the second one to the new place
                A[new_idx] = A[old_idx];
            }
            ++old_idx;
        }else{
            //If current element is not equal to previous one
            //Reset the count and keep the element
            ++new_idx;
            ++size;
            count = 1;
            A[new_idx] = A[old_idx];
            ++old_idx;
        }
    }
    return size;
}

Python Code:

# func: remove duplicates from sorted array that are at most twice
# @param A: input array A
# @return: length of the new array
# complexity: time O(n), space O(1)
def remove_duplicates(A):
    length = len(A)
    if length <= 2:
        return length
    count = 1
    new_idx = 0
    old_idx = 1
    new_length = 1
    while old_idx < length:
        if A[old_idx] == A[new_idx]:
            count += 1
            if count <= 2:
                new_idx += 1
                new_length += 1
                A[new_idx] = A[old_idx]
            old_idx += 1
        else:
            new_idx += 1
            new_length += 1
            A[new_idx] = A[old_idx]
            count = 1
            old_idx += 1

    return new_length