[프로그래머스 lv3] 모두 0으로 만들기 풀이

해설

주어지는 그래프는 트리이다. 

트리는 모두 연결되어있기 때문에 모든 정점의 합이 0이 되기만 한다면 규칙을 여러번 사용할지언정 모든 정점을 0으로 만들 수 있다.

따라서, 모든 정점의 합(sum)을 구하고 합이 0이 아닐 경우 -1을 리턴한다.

 

모든 정점의 합이 0이라면 이제는 최소의 규칙을 적용하여 모든 정점을 0으로 만들 방법을 찾아야 한다.

 

1) func(int cur)은 현재의 지점(cur)을 루트 노드로 하는 트리에 대하여, 자식 노드들의 숫자를 모두 상쇄시키고 남은 것을 a[cur]에 옮기는 함수이다. 말단에 있는 트리들 부터 0으로 만든 뒤, 남은 숫자들을 처리하여 루트 노드로 올린다. 밑에서 처리할 수 있는 것들은 미리미리 0으로 만들어야 최소한의 규칙을 사용할 수 있다.

 

2) 위의 사항을 구현하기 위하여 먼저, cur이 가지고 있는 자식 노드(x)들에 대해 func(x)를 수행한다. 이렇게 하면 각각의 자식 노드들을 루트 노드로 하는 트리에 대하여 상쇄될 수 있는 숫자들을 모두 없어지고 a[x]에 처리되지 못한 숫자만 남는다. 이때, 자식 노드(x)들과 현재 노드(cur)의 숫자를 상쇄하는데, 자식 노드를 0으로 만들어야 하므로 현재 노드에 a[x]를 더한다.

 

3) 상쇄를 처리할 때마다 해당 값만큼 cnt를 증가시킨다. 루트 노드까지 상쇄를 처리했을 때 cnt를 리턴하면 정답이다.

 

 

코드

#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <tuple>
#include <vector>
#include <stack>
#include <queue>
#include <deque>
#include <set>
#include <map>
#include <unordered_set>
#include <unordered_map>
#include <iostream>

using namespace std;
using ii = pair<int, int>;
using iii = tuple<int, int, int>;
using ll = long long;

#define X first
#define Y second

ll cnt;
vector<ll> aa(300005);
vector<bool> check(300005, false);
vector<vector<int>> graph(300005, vector<int>(0));

void func(int cur) {

    check[cur] = true;
    for (int x : graph[cur]) {
        if (check[x]) continue;
        func(x);
        cnt += abs(aa[x]);
        aa[cur] += aa[x];
    }   
}

void setGraph(const vector<int>& a, const vector<vector<int>>& edges) {

    for (int i=0; i<(int)a.size(); i++)
        aa[i] = 1LL*a[i];

    for (auto e : edges) {
        int u = e[0];
        int v = e[1];
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
}

long long solution(vector<int> a, vector<vector<int>> edges) {

    ll sum = 0;
    for (ll x : a) sum += x;
    if (sum) return -1; 

    setGraph(a, edges);
    func(0);    
    return cnt;
}