[프로그래머스/Python파이썬] 모두 0으로 만들기

문제 출처: https://programmers.co.kr/learn/courses/30/lessons/76503

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

 

 

문제 설명

DFS로 문제를 해결한다는 아이디어가 안 떠올라서 시간 초과가 났던 문제입니다. 정점을 모두 방문하는 방법으로 DFS가 왜 안 떠올랐을까요? 저는 바보입니다. 

 

만약 리스트 a의 원소의 합이 0이 되지 않는다면, 주어진 연산을 통해 트리의 모든 가중치를 0으로 만들 수 없으므로 -1을 반환합니다. 

 

DFS를 이용해서 정점을 탐색하는데, 제일 끝 정점(연결된 방문하지 않은 정점이 없는 경우)부터 0으로 만들어줍니다. 0으로 만들어준다는 것은 결국 자신을 dfs로 호출한 부모 정점에 자신의 가중치를 더하는 것과 같습니다. 현재 정점을 0으로 만들기 위한 연산의 횟수는 정점의 가중치의 절대값입니다. 

 

 

소스코드

import sys
sys.setrecursionlimit(300000)
from collections import defaultdict
answer = 0

def dfs(x, a, tree, visited):
    global answer
    visited[x] = 1

    for y in tree[x]:
        if not visited[y]:
            visited[y]
            a[x] += dfs(y, a, tree, visited)
    answer += abs(a[x])

    return a[x]


def solution(a, edges):
    global answer

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

    tree = defaultdict(list)
    for i, j in edges:
        tree[i].append(j)
        tree[j].append(i)

    visited = [0]*(len(a))

    dfs(0, a, tree, visited)

    return answer