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

문제 설명

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.

임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다. 하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다. 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요. (만약 처음부터 트리의 모든 정점의 가중치가 0이라면, 0을 return 해야 합니다.)

제한사항

• a의 길이는 2 이상 300,000 이하입니다.
• a의 모든 수는 각각 -1,000,000 이상 1,000,000 이하입니다.
• a[i]는 i번 정점의 가중치를 의미합니다.
• edges의 행의 개수는 (a의 길이 - 1)입니다.
• edges의 각 행은 [u, v] 2개의 정수로 이루어져 있으며, 이는 u번 정점과 v번 정점이 간선으로 연결되어 있음을 의미합니다.
• edges가 나타내는 그래프는 항상 트리로 주어집니다.

---

풀이

35분
리프 노드에서부터 0을 만들어 나가면서 루트 노드(0번 노드)에 값을 전파하는 DFS방식으로 풀이하였다. 값을 전파할 때 문제에서 주어지는 가중치 벡터를 그대로 사용하였는데, 전파 과정에서 int형 범위를 벗어나는 예외처리를 하지 않아서 처음에 틀렸었다.

코드

// 11:13 ~ 11:50
#include <string>
#include <vector>
#include <cmath>
using namespace std;


long long answer;
vector<int> visit;
vector<vector<int>> adj;
vector<vector<int>> edges;

long long dfs(int n,vector<long long>& a){
    long long local_sum = 0;
    
    for(int i =0; i < adj[n].size(); ++i){
        if(visit[adj[n][i]] == 0){
            visit[adj[n][i]] = 1;
            local_sum += dfs(adj[n][i], a);
        }
    }
    
    a[n] += local_sum;
    answer += abs(a[n]);
    
    return a[n];
}

long long solution(vector<int> a, vector<vector<int>> _edges) {
    adj.resize(a.size(), vector<int>());
    visit.resize(a.size(), 0);
    edges = _edges;
    vector<long long> a_long(a.begin(),a.end());
    
    answer = 0;
    
    for(vector<int> v : edges){
        adj[v[0]].push_back(v[1]);
        adj[v[1]].push_back(v[0]);
    }
    
    visit[0] = 1;
    long long res = dfs(0, a_long);
    
    
    if(res != 0) return -1;
    return answer;
}