프로그래머스 - 양과 늑대

문제

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

프로그래머스

 

 

풀이

루트 노드부터 인접 노드를 통해 갈 수 있는 모든 노드를 탐색한다.

여기선 BFS를 사용한다.

 

인접 노드를 이동하려면 양과 늑대의 수를 비교해야 한다.

아래 그림에서 0 -> 1로 바로 갈 순 없다.

0->2 방향으로 탐색을 진행한 후 양의 수가 더 많으면 1을 탐색할 수 있는데, 그 전까지 1번은 방문해도 3번과 4번 노드에 대해 BFS 작업을 진행할 수 없다.

트리의 노드를 구성할 때 아래와 같이 구성하면 위 정보를 한 번에 파악할 수 있다.

(노드 번호, 양의 수, 늑대의 수, 방문할 노드 집합)

 

방문할 노드 모음을(visit)를 집합으로 생성하여 1번, 3번 같은 늑대 노드는 적절한 양의 노드 수가 채워지기 전까지 visit 집합에서 빠지지 않도록 한다.

 

특정 노드를 방문하여 다음 노드를 큐에 집어넣는 작업이 되면 그 노드는 visit 집합에서 뺀다.

양 노드를 방문하면 양의 수를 + 1 해주고 늑대 노드를 만나면 늑대 수를 + 1 해준다.

 

방문 가능한 최대 양의 수는 BFS를 진행할 때마다 새로 업데이트 해준다.

 

 

코드

python

def solution(info, edges):
    answer = 0
 
    tree = [[] for _ in range(len(info))]
    for edge in edges:
        tree[edge[0]].append(edge[1])
 
    q = [(0, 1, 0, set())]  # 현재 위치, 양의 수, 늑대의 수, 방문한 노드 집합
 
    while q:
       current_node, sheep_count, wolf_count, visit = q.pop(0)
      answer = max(answer, sheep_count)
       visit.update(tree[current_node])
 
        for next_node in visit:
            if info[next_node]:
                if sheep_count > wolf_count + 1:
                    q.append((next_node, sheep_count, wolf_count + 1, visit - {next_node}))
            else:
                q.append((next_node, sheep_count + 1, wolf_count, visit - {next_node}))
 
    return answer