프로그래머스 - 중력 작용

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

코딩테스트 연습 - 중력 작용

 

 

시간초과 뜬다. 나중에 도전하겠다..

import sys
sys.setrecursionlimit(100000)

def solution(values, edges, queries):
    answer = []

    from collections import defaultdict
    G = defaultdict(list)
    rG = dict()
    sG = defaultdict(int) # 자식 합
    
    path = [0] * len(values)
    def init_sum(node):
        path[node - 1] = 1
        child_sum = values[node - 1]

        for next_node in G[node]:
            if path[next_node - 1] == 0:
                rG[next_node] = node
                child_sum += init_sum(next_node)
        
        sG[node] = child_sum
        
        return child_sum

    def query2(node, w):
        while rG.get(node) is not None:
            p_node = rG[node]
            sG[node] -= values[node - 1] - values[p_node - 1] # node 하위 sum 초기화
            values[node - 1], values[p_node - 1] = values[p_node - 1], values[node - 1] # get parent val
            node = p_node # 상위 노드로
        sG[1] -= values[0] - w
        values[0] = w

    sG[1] += values[0]
    for f, t in edges:
        G[f].append(t)
        G[t].append(f)
        
    init_sum(1)
    
    for n, q in queries:
        if q == -1: # query 1
            answer.append(sG[n])
        else: # query 2
            query2(n, q)

    return answer