[월간 코드 챌린지 시즌2]모두 0으로 만들기, 파이썬

문제

---

• 프로그래머스 모두 0으로 만들기 문제 링크

풀이

---

월간 코드 챌린지 시즌2 1, 2 번은 매우 쉬웠는데 3번에서 런타임에러가 나와서 4번은 풀지도 못했다. 3번이 레벨 3이기 때문에 3번까지는 무난히 풀 수 있어야 한다. 어쨌든 3번은 트리가 주어졌을 때 리프노드부터 탐색할 수 있는지 묻는 문제이다.

• 주어진 트리의 모든 노드 가중치 합이 0이 아니면 -1을 리턴한다, 0이라면 무조건 답을 구할 수 있다.
• 트리는 어떤 노드도 루트노드가 될 수 있기 때문에 시작점을 아무거나 잡아도 상관 없지만 길이가 2이상이라 편의상 0을 루트노드로 잡았다.
• 풀이 1의 경우 재귀를 사용했는데 파이썬은 기본 재귀 횟수를 1000으로 제한한다고 한다. 따라서 a의 최대 길이인 300,000으로 재귀 횟수 제한을 늘려주어야 런타임에러를 피할 수 있다.
• 재귀를 쓰기 싫다면 풀이 2처럼 스택을 활용하여 풀 수 있다. 원리는 같다.

코드

---

풀이 1 - dfs, 재귀 활용

from collections import defaultdict
import sys
sys.setrecursionlimit(300000)

def solution(a, edges):  
    if sum(a) != 0 :
        return -1
    answer = 0 
    dic = defaultdict(list)
    for edge in edges :
        u, v = edge[0], edge[1]
        dic[u].append(v)
        dic[v].append(u)
    answer = dfs(-1, 1, dic, a)
    return answer

def dfs(parent, child, dic, a) :
    ans = 0
    cand = dic[child]
    for nxt_child in cand :
        if nxt_child == parent : 
            continue
        ans += dfs(child, nxt_child, dic, a)
    if parent == -1 :
        return ans
    a[parent] += a[child]
    return ans + abs(a[child])

풀이 2 - dfs, stack 활용

from collections import defaultdict

def solution(a, edges):  
    if sum(a) != 0 :
        return -1
    answer = 0 
    dic = defaultdict(list)
    for edge in edges :
        u, v = edge[0], edge[1]
        dic[u].append(v)
        dic[v].append(u)
    check = [0] * len(a)
    stack = [[-1, 0]]
    while stack :
        parent, child = stack[-1][0], stack[-1][1]
        if check[child] :
            stack.pop()
            answer += abs(a[child])
            a[parent] += a[child]
        else :
            check[child] = 1
            for nxt_child in dic[child] :
                if not check[nxt_child] :
                    stack.append([child, nxt_child])
    return answer