문제
지난 번 월간 코딩 페스티벌에 출제되었던 모두 0으로 만들기 문제를 풀어보았습니다.
대회 당시에는 리프 노드부터 부모에게 값을 더해주는 방식으로 진행했었고, 트리 문제라는 강박에 무식하게 리프 노드를 삭제하면서 풀었었습니다. 제 기억으로는 45점에 나머지 테스트케이스는 시간초과가 났던 것으로 기억합니다.
만약 2차원 배열 형태로 그래프를 구성했다면 시간 초과를 받지 않았겠지만, 저는 dictionary로 그래프를 구성했기 때문에 리프노드를 직접 삭제하느라… ㅎㅎ… 시간초과가 났습니다. 심지어 리프노드 체크도, 어떤 식으로 했냐면.. 인접 노드를 1개만 갖는 노드를 리프 노드로 판단해, 계속 트리의 모든 노드를 순회하며 인접 노드가 1개뿐이라면 리프노드라고 판단하고, 이를 부모에게 전달해준 후 리프노드를 삭제했습니다. 삭제 자체(del)는 O(1)이나, 계속 노드를 순회하는 것이 문제였습니다.
풀이
dfs로 풀이합니다. 아무 노드나 root로 정한 후, 재귀를 돌면서 자식의 값을 스스로의 a 배열 값에 더하며 결과 값을 계속 부모 노드로 전달합니다.
그 모든 결과 값(root의 최종 값)이 0이 되면 성공입니다. 결과값(연산 시행 횟수)을 구하기 위해선, 어떠한 전역변수를 갖고, 재귀에서 탈출해 dfs를 빠져나가는 노드의 현재 a 배열 값을 더해줍니다. 단, 그 값이 음수일 수 있으므로 절댓값을 씌워서 받아주어야 합니다.
코드
import sys
sys.setrecursionlimit(10**9)
ret = 0
answer = 0
def dfs(start, a, graph, visited) :
global ret, answer
visited[start] = True
if start in graph :
for node in graph[start] :
if not visited[node] :
a[start] += dfs(node, a, graph, visited)
answer += abs(a[start])
return a[start]
def solution(a, edges):
if sum(a) != 0 : return -1
graph = {}
for edge in edges :
n1, n2 = edge[0], edge[1]
if n1 not in graph :
graph[n1] = [n2]
else :
graph[n1].append(n2)
if n2 not in graph :
graph[n2] = [n1]
else :
graph[n2].append(n1)
visited = [False]*len(a)
ret = dfs(0,a,graph,visited)
if ret != 0 : return -1
return answer