Iterating over vectors in C++ and erasing certain elements

Question

Just to provide background, I'm trying to write a function in C++ that takes a vector of ints and a particular int, and removes all of the elements of the vector in place that match the particular int, while returning the number of times it appears.

I'm not asking for help solving the problem. I am stuck trying to figure out why the wrong elements are being erased. Here is the code:

int removeElement(vector<int>& nums, int val) {
    int output = 0;
    int i = 0;
    while (i < nums.size()) {
        cout << nums[i] << "  " << i << "  " << (nums[i] == val) << "\n";
        if (nums[i] == val) {
            nums.erase(nums.begin() + i);
            output+=1;
        }
        else {
            i += 1;
        }
    }
    cout << "---------------\n";
    return output;
    }

Here is what I'm using to test it:

int main() {
    vector<int> firstOne = {3,2,2,3};
    cout << removeElement(firstOne,2) << "\n";
    firstOne = {3,2,2,3};
    cout << removeElement(firstOne,3) << "\n";
}

The output is supposed to be {3,3} and then {2,2} but it's {3,3} twice. Not only that, but the whole thing crashes when I try to test it with a vector with only 2 elements. I suspect there is a gap in my understanding of how vectors work. Can anyone explain to me what I'm doing wrong?

Answer 1

It is better to use erase-remove idiom

int removeElement(vector<int>& nums, int val) {
    int output = 0;
    int i = 0;
    // remove if moves elements "to be removed" in the end
    auto newend = std::remove_if(nums.begin(), nums.end(), [&](int element){
        cout << element << "  " << i << "  " << (element == val) << "\n";
        
        i++;
        if(element == val) {
            output++;
            return true; // if true, element will be removed
        }

        return false; // if false, element will not be removed
    });

    nums.erase(newend, nums.end());

    cout << "---------------\n";
    return output;
}

Answer 2

Besides the problems suggested in the comments that you just fixed, your code works fine. You are also not updating output anywhere so the function always returns 0 instead of the number of times the int appears in the passed vector . I edited your function a little and it's good to go:

int removeElement(vector<int>& nums, int val) {
    int output = 0;
    int i = 0;
    while (i < nums.size()) {
        cout << nums[i] << "  " << i << "  " << (nums[i] == val) << "\n";
        if (nums[i] == val) {
            nums.erase(nums.begin() + i);
            output++;
        }
        else {
            i += 1;
        }
    }
    cout << "---------------\n";
    return output;
}