Lesson 7: Stacks and Queues → Fish

문제

https://app.codility.com/programmers/lessons/7-stacks_and_queues/fish/

Fish coding task - Learn to Code - Codility

 

---

문제 풀이

서로 다른 방향으로 물고기가 만나는 경우, 크기가 큰 물고기가 작은 물고기를 잡아먹는다.

이를 위해 하류로 가는 물고기를 stack에 담고, 하류로 가는 물고기들과 만나면 크기를 비교하여 물고기를 제거한다.

같은 방향으로 향하는 물고기들은, 속도가 같기에 만날 수 없다는 제한 조건을 가지고 있다.

 

---

코드

def solution(A, B):
    stack = []
    cnt = 0

    for size, dir in zip(A, B):
        if dir == 1:
            stack.append(size)
        else:
            if not stack:
                cnt += 1
            else:
                while stack[-1] < size:
                    stack.pop()
                    if not stack:
                        cnt += 1
                        break

    return cnt + len(stack)