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

문제

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

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

풀이

각 leaf node에서 가중치를 부모 노드로 옮긴다.
이 과정을 leaf 노드가 사라질 때까지 반복한다.

 

시간초과

dfs 나 bfs를 이용해서 모든 노드를 탐색하도록 했으나, 시간초과로 대부분의 테스트 케이스에서 통과하지 못했다.
leaf node를 제거하면, 그 부모가 leaf node 후보가 되므로, 부모가 가진 간선개수를 이용해서 다음에 올 leaf node를 구할 수 있다. 매번 len으로 노드들의 간선 개수를 구하는 것도 시간이 오래 걸려서, 간선의 개수를 미리 구해서 저장해뒀다.

코드

from collections import defaultdict
def solution(a, edges):
    size = len(a)
    ans = 0
    
    # 노드 간 연결
    graph = defaultdict(set)
    for n, m in edges : 
        graph[n].add(m)
        graph[m].add(n)
    
    # 노드별 간선 개수
    count = dict()
    for n in range(size) : count[n] = len(graph[n])
    
    # 초기 leaf node
    leafnode = list()
    for n in range(size) : 
        if count[n]==1 : leafnode.append(n)
    
    while leafnode : 
        n = leafnode.pop()
        if count[n]==1 : 
            # disconnect
            pnode = graph[n].pop()       
            graph[pnode].remove(n)
            
            # update 간선 개수
            count[n] -= 1 
            count[pnode] -= 1
            
            # move weight
            ans += abs(a[n])
            a[pnode] += a[n]
            a[n] = 0
            
            # add next leafnode
            if count[pnode] == 1 : 
                leafnode.append(pnode)

    return -1 if any(a) else ans