[프로그래머스/Python] 양과 늑대

📌 문제 탐색하기

• 이진 트리 모양 초원의 각 노드에 늑대와 양이 한 마리씩 놓여 있다.
• 초원의 루트 노드에서 출발해 각 노드를 돌아다니며 양을 모으고자 한다.
  • 각 노드를 방문할 때마다 해당 노드의 양과 늑대가 따라온다.
  • 모은 양의 수 ≤ 모은 늑대의 수가 되면 늑대가 모든 양을 잡아먹는다.
• info = 각 노드에 있는 양 또는 늑대에 대한 정보가 담긴 배열
  • info의 원소는 0 또는 1
  • 2 ≤ info의 길이 ≤ 17
  • info[i]는 i번 노드에 있는 양 또는 늑대이며, info[0]은 항상 0
• edges = 이진 트리의 각 노드들의 연결 관계를 담은 2차원 배열
  • edges의 세로 길이 = info의 길이 -1
  • edges의 가로 길이 = 2
  • edges의 각 행 = [부모 노드 번호, 자식 노드 번호] 형태
  • 동일한 간선에 대한 정보가 중복해서 주어지지 않음
  • 0번 노드는 항상 루트 노드
• 최대한 많은 수의 양을 모아서 루트 노드로 돌아올 때, 몇 마리 모을 수 있는지 구하기

현재 노드를 기준으로 노드와 연결된 모든 노드로 한번씩 이동해서 결과를 갱신하면서 모든 경우의 수를 탐색하는, 백트래킹을 사용해 문제를 해결하고자 했다. 다만, 이미 지나온 노드로 다시 돌아가 다른 노드로 이동할 수 있으므로 이에 대한 처리에 주의했다.

 

📌 코드 설계하기

1. 주어진 info와 edges를 바탕으로 초원의 각 노드에 대한 정보를 저장할 grassland 리스트를 생성하고, 각 노드 요소마다 [해당 노드와 연결된 노드들]를 저장해 초기화한다. 
2. 모은 양의 최대 마리 수를 저장할 max_sheep 변수를 0으로 초기화한다.
3. 현재 노드의 방문 정보를 저장할 visited를 초기화한다.
   • visited[i][a][b] = i번째 노드에 양 a마리와 늑대 b마리를 가지고 방문한적 있는지 확인 
4. 백트래킹 알고리즘을 수행할 함수 backtracking(cur_node, cur_sheep, cur_wolf)를 정의한다.
   1. cur_node를 바탕으로  max_sheep를 갱신한다.
   2. 현재 노드와 연결된 노드를 하나씩 방문한다.
   3. 만약 연결된 노드의 visited를 확인해서 현재 (cur_node, cur_sheep)이 없는 경우에만 노드를 방문한다.
   4. 방문한 노드에 양 혹은 늑대가 있고, 여전히 양이 늑대보다 많다면
      1. info[방문한 노드]를 -1로 바꿔주고, cur_sheep과 cur_wolf를 갱신한다.
      2. 이 때, visited_way[방문한 노드][갱신된 cur_sheep][cur_wolf] 정보를 저장한다.
      3. backtracking(방문한 노드, 갱신된 cur_sheep, 갱신된 cur_wolf)를 재귀호출 한뒤, info와 visited를 복구해준다.
5. backtracking(0, 1, 0)를 호출해 갱신된 max_sheep를 반환한다.

 

📌 정답 코드

max_sheep = 0

def solution(info, edges):
    global max_sheep
    
    node_count = len(info)
    glassland = [[] for _ in range(node_count)]
    for a, b in edges:
        glassland[a].append(b)
        glassland[b].append(a)
    
    visited = [[[False]*18 for _ in range(18)] for _ in range(node_count)]
    
    def backtracking(cur_node, cur_sheep, cur_wolf):
        global max_sheep
        
        max_sheep = max(max_sheep, cur_sheep)
        for next_node in glassland[cur_node]:
            if not visited[next_node][cur_sheep][cur_wolf]:
                next_sheep, next_wolf = cur_sheep, cur_wolf
                next_node_info = info[next_node]
                if next_node_info == 1:
                    next_wolf += 1
                elif next_node_info == 0:
                    next_sheep += 1
                if next_sheep > next_wolf:
                    info[next_node] = -1
                    visited[next_node][next_sheep][next_wolf] = 1
                    backtracking(next_node, next_sheep, next_wolf)
                    info[next_node] = next_node_info
                    visited[next_node][next_sheep][next_wolf] = 0
    
    info[0] = -1
    visited[0][1][0] = 1
    backtracking(0, 1, 0)
    
    return max_sheep