(코딜리티) Lesson 7. [Stacks and Queues] : Fish - python

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

Fish coding task - Learn to Code - Codility

문제 내용

2개의 비어있지 않은 N개의 배열 A와 B가 주어질 때, A는 물고기의 크기를 담은 배열이고, B는 물고기의 방향을 담은 배열이다.

이 때, B는 0또는 1값만 가지며, 이는 아래와 같다.

• 0은 상류로 가는 물고기를 나타낸다.
• 1은 하류로 가는 물고기를 나타낸다.

만약 B[P] = 1이고, B[P] = 0이면 물고기는 만나며, 더 큰 사이즈 물고기가 잡아먹고 방향은 그대로 동일하다.

이때, 마지막으로 살아남은 물고기 수를 return하시오.

 

예를 들어,

A = [4, 3, 2, 1, 5], B = [0, 1, 0, 0, 0]일 경우

index 1 물고기를 제외한 나머지는 방향이 다르다. 우선 0번째 물고기는 살아남고, 2번, 3번은 1번물고기에 의해 잡아먹히며, 마지막 5번 물고기는 1번 물고기를 잡아먹어서 최종적으로 0번물고기와 5번물고기만 살아남기 때문에 2를 return한다.

제한 사항

• N는 1~100,000
• A의 각 요소는 0~1,000,000,000
• B의 각 요소는 0또는 1
• A의 각 요소는 중복되지 않는다.

코드

import sys
# 재귀함수를 사용하기 위한 limit를 N의 최대값으로 설정 
sys.setrecursionlimit(100000)

def solution(A, B):
    
    # 최종 물고기 수를 위한 배열
    fishes = []
    
    # 재귀 함수
    def stack_check():
        p = fishes.pop()
        if p[1] == 1 and ab[1] == 0:
            if p[0] > ab[0] :
                fishes.append(p)
            else :
                if fishes:
                    # fishes에 담겨있는 물고기를 계속 잡아먹을 수 있기 때문에 재귀함수로 호출
                    stack_check()
                else :
                    fishes.append(ab)
                
        else :
            fishes.append(p)
            fishes.append(ab)
	
    # A와 B를 합친 배열 for문
    for ab in zip(A,B):

        if not fishes:
            fishes.append(ab)
            continue

        stack_check()
        
    return len(fishes)

해설

코드가 좀 더러운 느낌이 없지 않은 듯한데,, while문이나 for문으로 풀 수 있을 것이라고 생각했으나, 재귀함수를 쓰고 싶어서 한번 사용해봤습니다. 

fishes에서 방향이 1,1,1,1,1 이다가 마지막 0이 제일 큰 물고기 사이즈인 경우는 fishes에 담겨있는 물고기를 다 잡아먹기 때문에, 이 경우에는 fishes배열을 잡아 먹지 않을 때 가지 반복문을 돌려줘야 합니다.