LeetCode: Remove Element

Given an array and a value, remove all instances of that value in place and return the new length.

The order of elements can be changed. It doesn't matter what you leave beyond the new length.

---

C++ Code:

/*
 * func: remove_element
 * goal: remove all elements of a given value in the array
 * @param A: integer array A
 * @param n: the size of the array
 * @param elem: target element
 */
int remove_element(int A[], int n, int elem){
    if(n <= 0){
        return n;
    }else{
        int i = 0;
        int j = n-1;
        while(i <= j){
            if(A[i] == elem){
                while(A[j] == elem){
                    --j;
                }
                if(i > j){
                    break;
                }
                std::swap(A[i], A[j]);
                --j;
            }else{
                ++i;
            }
        }
        return j+1;
    }
}

Python Code:

# func: remove all elements of a given value in the list
# @param A: the given list
# @param elem: target value
# @return: the new length of the list
def remove_list(A, elem):
    end = len(A) - 1
    for i in xrange(len(A)):
        if A[i] == elem:
            while end >= 0 and A[end] == elem:
                end -= 1
            if i >= end:
                break
            A[i] = A[end]
            A[end] = elem

    return end+1