Fish eat fish - linear solution

A bunch of cannibal fishes move up and downstream in a narrow river following the rules that I explained in the previous post. There I also showed a few test cases (C++ and GoogleTest) and a suboptimal first solution. Here I present a better algorithm that achieves an asymptotic linear time complexity even in the worst case.

The previously shown solution has the issue of repeating on and on checks on the same couple of fishes. We can avoid this delaying the comparison until it becomes a strict necessity.

The trick that I use here is avoid doing anything with any fish going downstream but storing them in a stack. The real job is done for each fish going upstream. If there is no one coming from the other direction, it survives for sure, otherwise it should decide its fate with a comparison against all the fishes previously stored in the stack, starting from the one closest to it (that's why a stack is a good choice as container).

Here is how I implemented this algorithm:

int solution(std::vector<int>& sizes, std::vector<int>& directions)
{
    int count = 0; // 1
    std::stack<int> goingDown; // 2
    for (unsigned i = 0; i < sizes.size(); ++i)
    {
        if (directions[i] == 1) // 3
        {
            goingDown.push(sizes[i]);
            continue;
        }

        if (goingDown.empty()) // 4
        {
            ++count;
            continue;
        }

        while (!goingDown.empty()) // 5
        {
            if (sizes[i] > goingDown.top()) // 6
                goingDown.pop();
            else
                break;
        }
        if (goingDown.empty()) // 7
            ++count;
    }
    return count + goingDown.size(); // 8
}

1. I don't need anymore to store the state of each fish, I could just keep track of the number of fishes going upstream that survived.
2. Here I push each fish that is going downstream.
3. The current fish is going downstream, I just have to push it in the stack and go to the next iteration.
4. The current fish is going upstream. If there is no previous fish moving in the other direction, he made it. I increase the counter and go to check the next one.
5. Otherwise, the current fish has to prove to be bigger than any fish in the stack to survive.
6. If the current fish is bigger that the current top in the stack, pop it (that is, eat it). Otherwise, it is the one to be eaten, and I should stop looping on the stack.
7. If the stack of downwards-going fishes is empty, this means that the current fish has eaten them all up (or it was so lucky to find no traffic). So I can increase the counter.
8. In "count" I have the number of all the upwards-going fishes that made it. I just have to add to it the numbers of the downwards-going fishes to get their current total number.

The C++ source code for both solutions and relative test cases is on github.