Solution to Fish task by codility

This task is a training on how to use stacks https://app.codility.com/programmers/lessons/7-stacks_and_queues/fish/

The two main important things to notice so you can complete the task are:

1. A fish going upstream will NEVER meet a downstream fish that comes next
2. A upstream fish can ether eat all the downstream fish or die

This is a possible solution using one stack:

def solution(A,B):
    alive = len(A)
    downstream = []
    for size,direction in zip(A,B):
        if direction == 1:
            downstream.append(size)
        else:
            while downstream:
                if size < downstream[-1]:
                    # upstream fish is eaten
                    alive = alive - 1
                    break
                else:
                    # upstream fish eats the other one
                    downstream.pop()
                    alive = alive - 1
    return alive