[프로그래머스] 모두 0으로 만들기

https://programmers.co.kr/learn/courses/30/lessons/76503

코딩테스트 연습 - 모두 0으로 만들기

 

문제를 이해하는 것부터 어려웠다.

 

1. sum(a) 가 0 이 아니라면 모든 가중치를 0 으로 만들 수 없으므로 -1 을 반환하면 된다. 반대로 0 이면 모든 점에서 루트 노드로의 이동이 가능하므로 0 을 만들 수 있다.

 

2. 트리구조에서 임의의 점이 루트 노드가 될 수 있으며 리프 노드부터 루트 노드로의 한 방향으로 모든 가중치를 0 으로 만들었을 때가 반드시 필요에 의한 최소한의 이동이므로 가장 작은 연산 횟수가 나온다. 편의상 루트 노드를 0 으로 잡았다.

 

3. DFS(깊이 우선 탐색) 을 이용한 재귀함수를 사용하여 처음으로 리프 노드의 a 배열의 원소 값을 바꾸고 answer 값의 연산 횟수를 더하여 최종적으로 루트 노드에서의 a[0] = 0 을 만들어 주고 answer 를 반환한다.

 

예제를 기준으로 (answer(연산 횟수), a[](가중치)) 라고 하자. DFS 로 아래의 순서로 연산이 이루어진다. 

- 리프 노드

  1 : [0, 0]

  2 : [0, 2]

  4 : [0, 2]

- 1 의 부모 노드(루트 노드)

  0 : [0, -5]

- 2, 4 의 부모 노드

  3 : [4(자식 노드들의 연산 횟수 + abs(자식 노드의 가중치) 들의 합), 5(현재 노드의 가중치 + 자식 노드들의 가중치의 합)]

- 3 의 부모 노드(루트 노드)

  0 : [9(자식 노드들의 연산 횟수 + abs(자식 노드의 가중치) 들의 합)), 0(현재 노드의 가중치 + 자식 노드들의 가중치의 합)]

 

4. 런타임 에러로 알아보니 파이썬에는 재귀의 깊이가 제한되어 있어서 setrecursionlimit 코드를 추가해준다. 

 

from collections import defaultdict
import sys

sys.setrecursionlimit(300000)
answer = 0

def solution(a, edges):
    global answer
    
    if sum(a) != 0:
        answer = -1
    else:
        dic = defaultdict(list)
        visit = [0] * len(a)

        for i in edges:
            dic[i[0]].append(i[1])
            dic[i[1]].append(i[0])

        rf(0, a, dic, visit)
        
    return answer

def rf(e, a, dic, visit):
    global answer
    visit[e] = 1
    
    for i in dic[e]:
        if visit[i] == 0: 
            a[e] = a[e] + rf(i, a, dic, visit)
    
    answer = answer + abs(a[e])
        
    return a[e]
                
print(solution([-5, 0, 2, 1, 2], [[0, 1], [3, 4], [2, 3], [0, 3]])) # 9