[프로그래머스 Level3] 모두 0으로 만들기 - 파이썬(Python) - 월간 코드 챌린지 시즌2

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

프로그래머스

문제 설명

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.

• 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다.

하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다. 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요. (만약 처음부터 트리의 모든 정점의 가중치가 0이라면, 0을 return 해야 합니다.)

 

 

---

가장 먼저, 모든 점들의 가중치를 0으로 만들기 위해서는 처음 모든 점들의 가중치들의 합이 0이 되어야하므로 이 부분은 미리 예외처리를 해준다.

if sum(a) != 0:
    return -1

 

주어진 edges를 이용하여 간선들을 연결한다.

connect = [[] for _ in range(len(a))]
for edge in edges:
    connect[edge[0]].append(edge[1])
    connect[edge[1]].append(edge[0])

 

DFS를 이용하여 연결된 간선을 하나씩 확인하며 가중치의 결과값을 구한다.

기존 점에 그 값을 더하고, 결과값에 그 크기를 더한다.

answer = 0

def dfs(i, a, connect, visited):
    global answer
    visited[i] = True
    for j in connect[i]:
        if not visited[j]:
            w = dfs(j, a, connect, visited)
            a[i] += w
            answer += abs(w)
    return a[i]

 

 

여기까지 했을 때 이상하게 런타임 에러가 계속 나서 확인해 보니, 다음과 같았다.

파이썬의 기본 재귀 깊이 제한은 1000으로 매우 얕은 편이다. 따라서 재귀로 문제를 풀 경우 드물지 않게 이 제한에 걸리게 되는데, 문제는 코딩테스트 환경에서는 에러 메시지를 볼 수 없다.

이 재귀 깊이 제한 문제를 해결하기 위해 코드의 상단에 다음과 같은 코드를 입력해주어야 한다.

import sys
sys.setrecursionlimit(300000)

 

 

 

<< 전체 코드 >>

import sys
sys.setrecursionlimit(300000)

answer = 0

def dfs(i, a, connect, visited):
    global answer
    visited[i] = True
    for j in connect[i]:
        if not visited[j]:
            w = dfs(j, a, connect, visited)
            a[i] += w
            answer += abs(w)
    return a[i]

def solution(a, edges):
    if sum(a) != 0:
        return -1
    
    connect = [[] for _ in range(len(a))]
    for edge in edges:
        connect[edge[0]].append(edge[1])
        connect[edge[1]].append(edge[0])
    
    visited = [False for _ in range(len(a))]
    dfs(0, a, connect, visited)
    return answer