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

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

프로그래머스

---

문제 설명

각 점에 가중치가 부여된 트리가 주어지고, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.
*임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고 다른 한 쪽은 1 감소시킵니다.

트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어질때,
트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 횟수를 return하시오.

제한사항 

2<= a의 길이 <= 300,000, a의 모든 수는 -1,000,000 ~ 1,000,000 사이입니다.
edges의 행의 개수는 (a의 길이 -1)이며 각 행은 [u, v] 2개의 정수로 이루어져 있으며, 이는 u번 정점과 v번 정점이 연결되어있음을 의미합니다.

---

코드 (1차 시도)

def solution(a, edges):
    answer = 0
    if sum(a) != 0:
        return -1
    
    degree = [0] * len(a)
    oneBranchGroup = set()
    tree = [set() for _ in range(len(a))]
    for start, end in edges:
        tree[start].add(end)
        tree[end].add(start)
        degree[start] += 1
        degree[end] += 1
    
    for branch in range(len(a)):
        if degree[branch] == 1:
            oneBranchGroup.add(branch)
            
    while len(oneBranchGroup) > 1:
        branch = oneBranchGroup.pop()
        next = tree[branch].pop()
        answer += abs(a[branch])
        a[next] += a[branch]
        tree[next].remove(branch)
        degree[next] -= 1
        if degree[next] == 1:
            oneBranchGroup.add(next)

    return answer

아이디어 

1. 트리 구조이고, 결국 모든 가중치가 0이 될려면 전체 가중치의 합이 0이어야 한다.
2. 각 가중치를 노드에서 다른 노드로 1번씩 이동해서 전부 모으면 최소 횟수일 것이다.
3. 바깥 리프노드에서 안쪽으로 모은다 생각하면 간선이 1개뿐인 노드에 있는 값을 연결된 노드에 옮기고, 그걸 반복하면 마지막에 중앙에 위치한 노드만 남게 된다. 
3-1. degree를 이용해 각 노드에 간선 개수를 정리하고, 간선이 1개뿐인 노드들을 모아 값을 이동하고, 다시 간선이 1개만 남은 노드들을 모아 반복한다.
결과 - 16/18. 실패 2개 시간초과
고민 - 위상정렬처럼 풀었을 때 시간을 더 빠르게 할 아이디어가 안떠오른다. 기존에도 set이 아닌 다른 자료구조로를 사용했다가 시간초과가 4개가 나와서 set으로 바꿔본 경우인데 이 방법에서 더 빠르게 바꿀만한 방법이 생각이 안난다...

 

코드 (2차 시도)

import sys
sys.setrecursionlimit(300001)
def dfs(node, a):
    global visited
    global answer
    global tree
    now = a[node]
    visited[node] = True
        
    for next in tree[node]:
        if not visited[next]:
            now += dfs(next, a)
    
    answer += abs(now)
    return now
    
def solution(a, edges):
    global visited
    global answer
    global tree
    
    if sum(a) != 0:
        return -1
    
    tree = [[] for _ in range(len(a))]
    for start, end in edges:
        tree[start].append(end)
        tree[end].append(start)
        
    visited = [False] * (len(a))
    answer = 0
    dfs(0, a)
    return answer

아이디어 

2차 시도

1. A -> B로 탐색하는걸 B->A로 옮긴거랑 동일한 과정이라 생각하고, 어느 시작점에서든 트리에 모든 부분을 끌어당긴다? 밀고 간다고 생각하면 탐색만 쭉 돌면서 1차 시도랑 동일하게 더해주고, 현재 위치에 옮겨온다.
2. 이걸 재귀 형식으로 어떻게 짜야할지가 생각이 안났고 짜서 통과 한 후 다른 분의 코드를 보고 조금 다듬은..? 아이디어를 짜고 실제로 돌아가게 구현하는 연습이 많이 필요한 것 같다.