[프로그래머스] 양과 늑대 - python

반응형

양과 늑대

문제

2진 트리 모양 초원의 각 노드에 늑대와 양이 한 마리씩 놓여 있습니다. 이 초원의 루트 노드에서 출발하여 각 노드를 돌아다니며 양을 모으려 합니다. 각 노드를 방문할 때 마다 해당 노드에 있던 양과 늑대가 당신을 따라오게 됩니다. 이때, 늑대는 양을 잡아먹을 기회를 노리고 있으며, 당신이 모은 양의 수보다 늑대의 수가 같거나 더 많아지면 바로 모든 양을 잡아먹어 버립니다. 당신은 중간에 양이 늑대에게 잡아먹히지 않도록 하면서 최대한 많은 수의 양을 모아서 다시 루트 노드로 돌아오려 합니다.

각 노드에 있는 양 또는 늑대에 대한 정보가 담긴 배열 info, 2진 트리의 각 노드들의 연결 관계를 담은 2차원 배열 edges가 매개변수로 주어질 때, 문제에 제시된 조건에 따라 각 노드를 방문하면서 모을 수 있는 양은 최대 몇 마리인지 return 하도록 solution 함수를 완성해주세요.

제한사항

• 2 ≤ info의 길이 ≤ 17
  • info의 원소는 0 또는 1 입니다.
  • info[i]는 i번 노드에 있는 양 또는 늑대를 나타냅니다.
  • 0은 양, 1은 늑대를 의미합니다.
  • info[0]의 값은 항상 0입니다. 즉, 0번 노드(루트 노드)에는 항상 양이 있습니다.
• edges의 세로(행) 길이 = info의 길이 - 1
  • edges의 가로(열) 길이 = 2
  • edges의 각 행은 [부모 노드 번호, 자식 노드 번호] 형태로, 서로 연결된 두 노드를 나타냅니다.
  • 동일한 간선에 대한 정보가 중복해서 주어지지 않습니다.
  • 항상 하나의 이진 트리 형태로 입력이 주어지며, 잘못된 데이터가 주어지는 경우는 없습니다.
  • 0번 노드는 항상 루트 노드입니다.

풀이

제한사항을 보면 완전 탐색 백트래킹임을 알 수 있다. 자식 노드를 방문했다가 다시 부모 노드로 돌아와야하는 것 같아서 헷갈렸는데 그럴 필요가 없었다. 이상하게 안 풀려서 한참을 붙잡고 있었는데 막상 풀리고 나니 간단했던 문제 ㅠ

질문하기에 아주 좋은 설명글이 있어서 참고했다.

그래프 구성하기

단방향으로 인접리스트를 만들고 엣지 정보를 저장했다.

이미 지나갔던 정점은 다시 방문해도 아무런 영향이 없다. 최초 방문시에 정점에 있던 동물이 따라오기 때문이다. 따라서 이미 방문한 정점을 다시 방문하는 것이 아니라 다음 차례에 이동 가능한 정점을 따로 담아야 한다.

1번 예시의 경우 0번을 방문하면 [1, 8]의 이동 가능한 정점 가질 수 있다. 이런 식으로 부모 노드를 방문할 경우 그 자식 노드 역시 방문할 수 있기 때문에 방문할 수 있는 자식 노드들로 이동 가능한 정점을 업데이트해주면 된다. 이렇게 하면 그래프를 역류하지 않아도 되고 단방향 그래프로 구성해도 모든 노드를 방문할 수 있게 된다.

DFS 수행하기

DFS 함수의 인수를 4개로 정했다. (현재 정점, 양의 수, 늑대의 수, 이동 가능한 정점)

1. 현재 정점의 info 값을 확인한다.
2. 양일 경우 양의 수에 +1 해주고 양의 수가 최대인지 비교해서 업데이트한다.
3. 늑대일 경우 늑대의 수를 늘려준다.
4. 늑대가 양보다 같거나 많을 경우 종료한다. (리턴한다)
5. 현재 정점의 자식 노드를 이동 가능한 정점에 넣어준다.
6. 이동 가능한 정점을 탐색하면서 DFS를 진행한다. 해당 정점을 선택할 경우 방문한 정점이 되기 때문에 이동 가능한 정점에서는 제거하고 DFS 함수에 넣어준다.

전체 코드

def dfs(idx, sheep, wolf, possible):
    global g_info, answer, graph
    if g_info[idx] == 0:
        sheep += 1
        answer = max(answer, sheep)
    else:
        wolf += 1
        
    if wolf >= sheep:
        return 
    
    possible.extend(graph[idx])
    for p in possible:
        dfs(p, sheep, wolf, [i for i in possible if i != p])
    
def solution(info, edges):
    global answer, g_info, visited, graph
    answer = 0
    g_info = info
    n = len(info)
    graph = [[] for _ in range(n)]
    
    for a, b in edges:
        graph[a].append(b)
        
    
    dfs(0, 0, 0, [])
    return answer