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

이 문제는 dfs를 이용해서 풀었다.

 

이 문제에서 가장 핵심은 해당 그래프가 트리라는 것이다. 하나는 +1, 하나는 -1 이라는 것은 -1되는 쪽에서 +1되는 쪽으로 1을 준다라고 생각하면 쉽게 문제가 풀린다.

 

주의할 점은 단순히 a를 순회하면서 합이 0이 되지 않았다고 바로 -1을 리턴하는 경우 정답이 나오지 않는다는 것이다.

dfs 결과를 통해 총 합이 0이 되지 않을 때 -1을 리턴해야 정답이 나온다.

#include <string>
#include <vector>
#include <queue>
#include <algorithm>
#include <cmath>
#include <map>
using namespace std;

#define NODE_MAX 300000
#define INF 10000000000000000ll
typedef long long LL;

vector<int> edges[NODE_MAX];
vector<int> roots;

bool isAllZero(vector<int> a)
{
    for (int number : a)
    {
        if (number != 0)
        {
            return false;
        }
    }
    return true;
}

bool isPossible(vector<int> a)
{
    LL result = 0;
    for (int number : a)
    {
        result += number;
    }
    if (result != 0)
    {
        return false;
    }
    return true;
}

LL cnt;
vector<bool> visited;

LL dfs(int num, vector<int>& a)
{
    visited[num] = true;
    
    LL result = a[num];
    for (int next : edges[num])
    {
        if (!visited[next])
        {
            result += dfs(next, a);
        }
    }
    
    cnt += abs(result);
    return result;
}

LL solution(vector<int> a, vector<vector<int>> eds) {
    if (isAllZero(a))
    {
        return 0;
    }
    
    int size = a.size();
    for (vector<int>& ed : eds)
    {
        int u = ed[0];
        int v = ed[1];
        edges[u].push_back(v);
        edges[v].push_back(u);
    }
    
    visited = vector<bool>(size, false);
    cnt = 0;
    LL result = dfs(0, a);
    
    if (result != 0)
    {
        return -1;
    }
    
    return cnt;
}