문제 설명
2진 트리 모양 초원의 각 노드에 늑대와 양이 한 마리씩 놓여 있습니다. 이 초원의 루트 노드에서 출발하여 각 노드를 돌아다니며 양을 모으려 합니다. 각 노드를 방문할 때 마다 해당 노드에 있던 양과 늑대가 당신을 따라오게 됩니다. 이때, 늑대는 양을 잡아먹을 기회를 노리고 있으며, 당신이 모은 양의 수보다 늑대의 수가 같거나 더 많아지면 바로 모든 양을 잡아먹어 버립니다. 당신은 중간에 양이 늑대에게 잡아먹히지 않도록 하면서 최대한 많은 수의 양을 모아서 다시 루트 노드로 돌아오려 합니다.
예를 들어, 위 그림의 경우(루트 노드에는 항상 양이 있습니다) 0번 노드(루트 노드)에서 출발하면 양을 한마리 모을 수 있습니다. 다음으로 1번 노드로 이동하면 당신이 모은 양은 두 마리가 됩니다. 이때, 바로 4번 노드로 이동하면 늑대 한 마리가 당신을 따라오게 됩니다. 아직은 양 2마리, 늑대 1마리로 양이 잡아먹히지 않지만, 이후에 갈 수 있는 아직 방문하지 않은 모든 노드(2, 3, 6, 8번)에는 늑대가 있습니다. 이어서 늑대가 있는 노드로 이동한다면(예를 들어 바로 6번 노드로 이동한다면) 양 2마리, 늑대 2마리가 되어 양이 모두 잡아먹힙니다. 여기서는 0번, 1번 노드를 방문하여 양을 2마리 모은 후, 8번 노드로 이동한 후(양 2마리 늑대 1마리) 이어서 7번, 9번 노드를 방문하면 양 4마리 늑대 1마리가 됩니다. 이제 4번, 6번 노드로 이동하면 양 4마리, 늑대 3마리가 되며, 이제 5번 노드로 이동할 수 있게 됩니다. 따라서 양을 최대 5마리 모을 수 있습니다.
각 노드에 있는 양 또는 늑대에 대한 정보가 담긴 배열 info, 2진 트리의 각 노드들의 연결 관계를 담은 2차원 배열 edges가 매개변수로 주어질 때, 문제에 제시된 조건에 따라 각 노드를 방문하면서 모을 수 있는 양은 최대 몇 마리인지 return 하도록 solution 함수를 완성해주세요.
제한 사항
- 2 ≤
info의 길이 ≤ 17info의 원소는 0 또는 1 입니다.- info[i]는 i번 노드에 있는 양 또는 늑대를 나타냅니다.
- 0은 양, 1은 늑대를 의미합니다.
- info[0]의 값은 항상 0입니다. 즉, 0번 노드(루트 노드)에는 항상 양이 있습니다.
edges의 세로(행) 길이 =info의 길이 - 1edges의 가로(열) 길이 = 2edges의 각 행은 [부모 노드 번호, 자식 노드 번호] 형태로, 서로 연결된 두 노드를 나타냅니다.- 동일한 간선에 대한 정보가 중복해서 주어지지 않습니다.
- 항상 하나의 이진 트리 형태로 입력이 주어지며, 잘못된 데이터가 주어지는 경우는 없습니다.
- 0번 노드는 항상 루트 노드입니다.
코드
실패 코드 - 첫 시도
def solution(info, edges):
answer = []
graph = [[] for _ in range(len(info))] # 연결 정보
check = [0] * len(info) # 방문 처리 리스트
# 연결 정보 초기화
for p, c in edges:
graph[p].append(c)
def dfs(p, sheep, wolf):
if sheep > wolf: # 양의 수가 늑대의 수보다 많다면
answer.append(sheep)
else: # 늑대의 수가 더 많다면, 종료
return
# 연결된 자식 노드 탐색
for nxt in graph[p]:
if not check[nxt]: # 자식 노드를 방문하지 않았을 경우
check[nxt] = 1 # 방문 처리
if info[nxt] == 0: # 양이라면
dfs(nxt, sheep + 1, wolf)
else: # 늑대라면
dfs(nxt, sheep, wolf + 1)
check[nxt] = 0 # 방문 해제
dfs(0, 1, 0)
return max(answer)아이디어 & 해결
문제의 예시를 보면, 루트노드에서 자식노드로 갔다가 형제노드로 가는 방법도 가능하다.
하지만 위의 코드는, 자식노드에서 형제노드로 가는 경우의 수를 고려하지 않는다.
즉, 노드 p에 연결된 자식노드인 nxt를 재귀호출하고 다시 nxt의 자식노드를 재귀호출하는 방식이다.
따라서, 우리는 위의 코드에 형제노드로 갈 수 있는 코드를 추가해야한다.
성공 코드-#1
_max = 1
def solution(info, edges):
graph = [[] for _ in range(len(info))] # 연결 정보
# 연결 정보 초기화
for p, c in edges:
graph[p].append(c)
# DFS
# 현재 정점, 양의 수, 늑대의 수, 이동 가능한 정점 리스트
def dfs(p, sheep, wolf, node_list):
global _max
# 양의 수가 더 많으면
if sheep > wolf:
_max = max(_max, sheep) # 최대 양의 수 갱신
else: # 늑대의 수가 더 많으면
return
# 이동 가능한 정점 리스트 갱신
node_list.extend(graph[p])
for node in node_list:
# 자기 자신이 아닌 정점 리스트로 호출
if info[node] == 0: # 양이라면
dfs(node, sheep + 1, wolf, [i for i in node_list if i != node])
else: # 늑대라면
dfs(node, sheep, wolf + 1, [i for i in node_list if i != node])
# DFS 호출
dfs(0, 1, 0, [])
return _max아이디어 & 해결
위 코드는 형제노드로 가는 경우의 수를 추가했다.
node_list는 현재 노드 p에 연결된 자식 노드들을 담고있다.
node_list안의 노드들을 탐색하면서, 양인지 늑대인지에 따라 경우의 수를 나눈다.
[i for i in node_list if i != node]를 node_list의 인자로 설정하여 재귀호출을 하는데,
이 코드의 의미는 다음과 같다.
node를 재귀호출하므로 자기 자신(node)은 빼라는 것이다.
만약 자기 자신도 node_list에 추가한다면 무한루프에 빠질 가능성이 있다.
성공 코드-#2
def solution(info, edges):
answer = [] # 모을 수 있는 '양의 개수'
check = [0] * len(info) # 0-index, 방문 처리 리스트
check[0] = 1 # 루트 노드 방문 처리
def dfs(sheep, wolf):
# 늑대 수가 양의 수보다 같거나 많다면
# 양이 모두 잡아먹힌다.(양 : 0마리)
# 즉, answer에 넣지 않아도 된다.
if sheep <= wolf:
return
else: # 양의 수가 늑대의 수보다 '더 많은' 경우
answer.append(sheep)
# 연결 정보를 확인한다.
for parent, child in edges:
# 부모 노드는 무조건 방문을 해야 자식 노드를 방문할 수 있다.
# 자식 노드를 이미 방문한 경우, 재방문은 안된다.
if check[parent] and not check[child]:
check[child] = 1 # 자식노드 방문 처리
if info[child] == 0: # '양'이라면?
dfs(sheep + 1, wolf)
else: # '늑대'라면?
dfs(sheep, wolf + 1)
# 다음 경우의 수를 위해
# 자식노드 방문 해제
check[child] = 0
dfs(1, 0) # 루트 노드는 항상 양이다.
return max(answer)아이디어 & 해결
문제의 제한사항을 살펴보면 다음의 문구를 볼 수 있다.
✅ edges는 항상 하나의 이진 트리 형태로 입력이 주어진다.
이 말은 edges를 가지고 굳이 연결 정보를 다시 만들 필요는 없다는 것이다.
즉, 우리는 주어진 edges를 통해서 DFS를 수행하면 된다.
여기서 키 포인트는,
트리를 방문할 때 어떤 자식 노드를 방문하려고 한다면, 그 자식노드의 부모노드는 이미 방문했어야 한다는 것이다.
그 조건을 의미하는 것이 위의 코드에서 if check[parent] and not check[child]이다.
나머지 코드는 성공 코드-#1과 유사하다.
