[Programmers] 양과 늑대 (Python 풀이) - 2022 KAKAO BLIND RECRUITMENT

https://programmers.co.kr/learn/courses/30/lessons/92343

코딩테스트 연습 - 양과 늑대

---

풀이

이 문제는 이진 트리에서 조건에 따라 순회 가능한 모든 경우의 수를 탐색하면서 최대 양의 개수를 구하면 정답이 됩니다. 순회는 순서에 따라 영향을 받습니다. 따라서 같은 번호를 순회하더라도 여러 경우의 수가 생길 수 있습니다.

입력 및 초기화

from collections import defaultdict, deque  #
from copy import deepcopy
 
is_wolf = list()
num2edges = defaultdict(list)
max_sheep = 0
 
def solution(info, edges):
    global is_wolf, num2edges, max_sheep
    is_wolf = info # 전역변수에 할당
    used = [False] * len(is_wolf) # 방문한 노드 확인 용 변수
 
    for e_from, e_to in edges:
        num2edges[e_from].append(e_to) # 연결된 엣지 리스트에 추가
 
    # start
    find_max_recursive(0, used, 0, 0, [])
    return max_sheep

 

10번 줄에서 info는 is_wolf에 할당합니다. 함수에서 참고하기 위해 전역변수로 설정하였습니다.

13~14번 줄에서 num2edges는 인덱스로 정해진 번호에 연결된 노드들을 리스트로 가지는 변수입니다. 

위와 같이 노드가 구성되어 있을 때 numn2edges 값을 다음과 같습니다. 출력해보면 다음과 같은 값을 가집니다.

# num2edges
{
    0: [1, 8], 
    1: [2, 4], 
    8: [7, 9], 
    9: [10, 11], 
    4: [3, 6],
    6: [5]
}

입력 및 초기화가 완료되면 find_max_recursive() 함수에서 최대 양의 수를 찾습니다.

최대 양의 수 찾기

find_max_recursive() 함수는 다음과 같이 재귀함수로 이루어져 있습니다.

def find_max_recursive(current_loc, used, nsheep, nwolf, can_go):
    global max_sheep
 
    if used[current_loc]: return  # 현재 노드를 방문한 경우 return
    used[current_loc] = True  # 방문 기록
 
    if is_wolf[current_loc]:  # 늑대인 경우 늑대 count += 1
        nwolf += 1
    else:
        nsheep += 1  # 양인 경우 양 count += 1
        max_sheep = max(max_sheep, nsheep)  # 양 최대 수 갱신
 
    if nsheep <= nwolf: return  # 늑대 수가 양의 수와 같거나 많은 경우 return
 
    can_go.extend(num2edges[current_loc])  # 현재 노드와 연결된 노드를 추가함
    for next_loc in can_go:
        # q에 저장된 노드서 하나를 가져와서 재귀함수 요청
        # 이 때 다음 q에는 현재 노드를 제외한 리스트로 재구성하여 재귀함수 요청
        find_max_recursive(next_loc, deepcopy(used), nsheep, nwolf,
                           can_go=[loc for loc in can_go if loc != next_loc and not used[loc]])

 

7~11번 줄에서 현재 노드가 양인지 늑대인지에 따라 양의 수 혹은 늑대의 수를 갱신합니다. 양의 수가 더해진 경우 양 최대 수를 갱신합니다.

13번 줄에서 늑대 수가 양의 수보다 많아진 경우 종료합니다.

15번 줄에서 현재 노드와 연결된 노드들을 can_go 리스트에 추가합니다.

16~19번 줄에서 can_go 리스트에 저장된 노드들을 하나씩 가져와서 재귀함수를 요청합니다. 다음 재귀함수에서 current_loc 가 next_loc가 됩니다. 그리고 can_go 리스트에 저장된 노드들 중 next_loc는 사용한 변수가 되기 때문에 next_loc를 제외한 노드들로 can_go를 재구성하여 재귀함수를 요청합니다.

전체 소스 코드

from collections import defaultdict, deque  #
from copy import deepcopy
 
is_wolf = list()
num2edges = defaultdict(list)
max_sheep = 0
 
 
def find_max_recursive(current_loc, used, nsheep, nwolf, can_go):
    global max_sheep
 
    if used[current_loc]: return  # 현재 노드를 방문한 경우 return
    used[current_loc] = True  # 방문 기록
 
    if is_wolf[current_loc]:  # 늑대인 경우 늑대 count += 1
        nwolf += 1
    else:
        nsheep += 1  # 양인 경우 양 count += 1
        max_sheep = max(max_sheep, nsheep)  # 양 최대 수 갱신
 
    if nsheep <= nwolf: return  # 늑대 수가 양의 수와 같거나 많은 경우 return
 
    can_go.extend(num2edges[current_loc])  # 현재 노드와 연결된 노드를 추가함
    for next_loc in can_go:
        # q에 저장된 노드서 하나를 가져와서 재귀함수 요청
        # 이 때 다음 q에는 현재 노드를 제외한 리스트로 재구성하여 재귀함수 요청
        find_max_recursive(next_loc, deepcopy(used), nsheep, nwolf,
                           can_go=[loc for loc in can_go if loc != next_loc and not used[loc]])
 
 
def solution(info, edges):
    global is_wolf, num2edges, max_sheep
    is_wolf = info # 전역변수에 할당
    used = [False] * len(is_wolf) # 방문한 노드 확인 용 변수
 
    for e_from, e_to in edges:
        num2edges[e_from].append(e_to) # 연결된 엣지 리스트에 추가
    print(num2edges)
    for i, v in num2edges.items():
        print(i, ":", v)
    # start
    find_max_recursive(0, used, 0, 0, [])
    return max_sheep