Level 3 모두 0으로 만들기 Python 3

https://programmers.co.kr/learn/courses/30/lessons/76503?language=python3 

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

 

최종 코드

• GitHub -> https://github.com/hwayeon351/Programmers-Algorithms/blob/main/level3_모두0으로만들기.py

hwayeon351/Programmers-Algorithms

import sys
from collections import defaultdict
sys.setrecursionlimit(300000)
def dfs(now_i, before_i, visit, adj, a):
    global answer
    for next_i in adj[now_i]:
        if not visit[next_i]:
            visit[next_i] = True
            dfs(next_i, now_i, visit, adj, a)
    a[before_i] += a[now_i]
    answer += abs(a[now_i])
    a[now_i] = 0
      
def solution(a, edges):
    global answer
    answer = 0
    if sum(a)!=0: return -1
    if a.count(0) == len(a): return 0
    adj = defaultdict(set)
    for n1, n2 in edges:
        adj[n1].add(n2)
        adj[n2].add(n1)
    visit = [False]*len(a)
    dfs(0, 0, visit, adj, a)
    return answer

---

풀이 과정

처음에 문제를 접근했을 때 트리의 단말 노드로부터 루트까지 가중치를 올려나가는 방법으로 구현을 하려고 했다.

그러기 위해서는 특정 노드의 부모가 누구인지와 특정 노드의 레벨을 표현해줘야 했다.

입력인 edges의 항목을 하나씩 받아오면서 트리를 만들어가면서 부모와 레벨을 체크해줬는데 복잡했다.

 

사실상 항상 트리 구조로 입력이 들어오기 때문에 어디서부터 시작하는지는 중요하지 않았다. 모든 노드를 탐색해서 한 노드로 가중치 값을 모으면 되는 문제였다.

따라서, 특정 노드로부터 단말 노드까지 탐색하면서 단말 노드에 도착한 경우 부모 노드에게 가중치 값을 주고 카운팅하는 방식으로 문제를 해결했다. 깊이 우선 탐색인 dfs 알고리즘을 적용하였다.

import sys
from collections import defaultdict
#a의 길이가 최대 300000이고, 최대 edges 개수는 299999이므로
sys.setrecursionlimit(300000)
def dfs(now_i, before_i, visit, adj, a):
    global answer
    #현재 노드에 인접한 노드들을 모두 탐색한다 -> 깊이 우선 탐색
    for next_i in adj[now_i]:
        if not visit[next_i]:
            visit[next_i] = True
            dfs(next_i, now_i, visit, adj, a)
    #인접한 노드들을 모두 완료한 후 
    #이전에 방문한 노드의 가중치에 현재 노드의 가중치를 더함
    a[before_i] += a[now_i]
    #|현재 노드의 가중치|를 answer에 더함 
    answer += abs(a[now_i])
    #현재 노드의 가중치를 0으로 바꿔줌
    a[now_i] = 0
      
def solution(a, edges):
    global answer
    answer = 0
    #모든 노드의 가중치의 합이 0이 아니면 가중치를 0으로 만들 수 없는 경우
    if sum(a)!=0: return -1
    #처음부터 모든 가중치가 0인 경우
    if a.count(0) == len(a): return 0
    adj = defaultdict(set)
    for n1, n2 in edges:
        adj[n1].add(n2)
        adj[n2].add(n1)
    visit = [False]*len(a)
    dfs(0, 0, visit, adj, a)
    return answer
#시간복잡도 = O(V), 공간복잡도 = O(V)

DFS의 시간복잡도는 O(V+E)이며 이 문제에서는 항상 트리 구조의 그래프를 탐색하므로 E = V-1이다. 따라서, 시간복잡도는 O(2V-1) = O(V)이다. 공간복잡도는 인접리스트를 구성하는데에서 발생하는 O(2(V-1)) = O(V)이다.