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

 

트리의 특정 간선에 연결된 두 노드 중, 한 쪽의 가중치에는 -1을 하고 다른 쪽의 가중치에는 +1을 하는 작업을 반복했을 때, 모든 노드의 가중치를 0으로 만들 수 있을까?

 

만약, 만들 수 없다면 -1을 출력하면 되고 만들 수 있다면 몇 번을 반복해야 모든 노드의 가중치가 0이 되는지 구하면 된다.

 

문제 풀이

언뜻 보기엔, 매우 복잡해보이지만 생각보다는 간단한 문제이다.

핵심은 트리의 가장 말단 노드부터 0으로 만드는 것이다.

 

DFS든 BFS든 트리의 말단을 찾아가서, 말단을 0으로 만들면서 루트 노드까지 올라왔을 때 모든 노드의 가중치를 확인하면 된다.

 

그림을 보면서 이해해보자.

 

문제의 첫 번째 입력을 그림으로 그려보면 아래와 같다.

노드 안의 숫자는 (번호 / 가중치) 이다.

이제, 말단에서부터 가중치를 0으로 만들어보자.

 

2번 노드에 대해 갱신하면, 위와 같아진다.

다음은 4번 노드에 대해서도 갱신해보자.

위와 같이 값이 갱신될 것이다.

이제, 말단 노드가 모두 0으로 되었으니, 0이 아닌 노드 중 가장 말단에 있는 노드를 0으로 만들어보자.

 

3번 노드를 갱신하고 나니, 루트 노드까지 갱신되었다.

노드의 모든 가중치를 보면 0으로 갱신되어 있는 것을 볼 수 있다.

 

이처럼, 말단노드부터 위로 올라오며 0으로 갱신하다 보면, 모든 노드를 갱신하게 된다.

모든 노드를 갱신한 뒤에, 모든 노드의 가중치가 0인지 아닌지를 확인하면 된다.

 

가중치가 0이 아니라면, -1과 +1을 하는 작업을 몇 번 반복했는지를 출력하면 되는데, 이건 노드를 0으로 만들 때, 해당 노드의 가중치의 절대값에 대한 누적합을 구하면 된다.

 

참고로 위의 과정은 루트 노드에서부터 시작해야 한다.

하지만, 입력으로 주어지는 트리는 양방향 트리이므로, 어떤 노드든 사실 루트 노드가 될 수 있다.

그러므로 아무 노드에서나 시작해도 문제는 없다.

 

하지만, 4번 노드에서부터 시작하도록 했는데 입력으로 4번 노드가 주어지지 않는다면?

이런 문제를 방지하기 위해, 0번 노드에서부터 시작하는 것이 가장 무난하다. 

 

풀이 코드

std::vector<long long> CastA(a.size());
for(int i = 0; i < a.size(); i++)
{
    CastA[i] = a[i];
}

std::vector<std::vector<int>> Link(a.size());

for(int i = 0; i < edges.size(); i++)
{
    int Node_1 = edges[i][0];
    int Node_2 = edges[i][1];

    Link[Node_1].push_back(Node_2);
    Link[Node_2].push_back(Node_1);
}
    
std::vector<bool> isVisit(a.size(), false);

CastA는 주어진 노드의 가중치를 long long 자료형으로 변환하여 새로 저장한 것이다.

왜냐하면, 기존 배열에 값을 더하고 빼며 계산할 것인데 이 과정에서 int 자료형의 범위를 초과하는 경우가 발생할 수 있기 때문이다.

 

Link는 각 노드에서 연결된 노드들을 저장한 배열이다.

 

isVisit는 DFS를 위해 만든 방문체크배열이다.

 

DFS(CastA, Link, isVisit, 0, -1);

for(int i = 0; i < CastA.size(); i++)
{
    if(CastA[i] != 0)
    {
        Count = -1;
        break;
    }
}

return Count;

다음은 DFS를 돌린 다음, 모든 노드의 가중치가 0인지를 확인하는 과정이다.

(참고로 Count는 전역변수이다.)

 

void DFS(std::vector<long long>& _Weight, const std::vector<std::vector<int>>& _Link,  std::vector<bool>& _isVisit, int _CurIndex, int _PrevIndex)
{
    _isVisit[_CurIndex] = true;
    
    for(int i = 0; i < _Link[_CurIndex].size(); i++)
    {
        int NextIndex = _Link[_CurIndex][i];
        
        if(_isVisit[NextIndex] == false)
        {
            DFS(_Weight, _Link, _isVisit, NextIndex, _CurIndex);
        }
    }
    
    if(_PrevIndex != -1)
    {
        Count += abs(_Weight[_CurIndex]);

        _Weight[_PrevIndex] += _Weight[_CurIndex];
        _Weight[_CurIndex] = 0;
    }
}

DFS내부 코드는 이와 같다.

 

먼저, 방문체크를 해준 다음 연결된 노드 중 이동할 수 있는 노드에 대해 이동해주었다.

이후, 말단 노드에서는 더이상 이동할 곳이 없기 때문에 반복문을 탈출하게 되는데, 이 때, 해당 노드의 가중치를 0으로 만드는 과정을 거치게 된다. 말단 노드에서부터 위로 올라오며 재귀적으로 노드의 가중치를 0으로 만드는 것이다.

 

 이렇게 하면, 올바른 답을 출력할 수 있게 된다.

 

 

코드 전문

#include <string>
#include <vector>

using namespace std;

long long Count = 0;

void DFS(std::vector<long long>& _Weight, const std::vector<std::vector<int>>& _Link,  std::vector<bool>& _isVisit, int _CurIndex, int _PrevIndex)
{
    _isVisit[_CurIndex] = true;
    
    for(int i = 0; i < _Link[_CurIndex].size(); i++)
    {
        int NextIndex = _Link[_CurIndex][i];
        
        if(_isVisit[NextIndex] == false)
        {
            DFS(_Weight, _Link, _isVisit, NextIndex, _CurIndex);
        }
    }
    
    if(_PrevIndex != -1)
    {
        Count += abs(_Weight[_CurIndex]);

        _Weight[_PrevIndex] += _Weight[_CurIndex];
        _Weight[_CurIndex] = 0;
    }
}

long long solution(vector<int> a, vector<vector<int>> edges) 
{
    std::vector<long long> CastA(a.size());
    for(int i = 0; i < a.size(); i++)
    {
        CastA[i] = a[i];
    }
    
    std::vector<std::vector<int>> Link(a.size());
    
    for(int i = 0; i < edges.size(); i++)
    {
        int Node_1 = edges[i][0];
        int Node_2 = edges[i][1];
        
        Link[Node_1].push_back(Node_2);
        Link[Node_2].push_back(Node_1);
    }
    
    std::vector<bool> isVisit(a.size(), false);
    
    DFS(CastA, Link, isVisit, 0, -1);
    
    for(int i = 0; i < CastA.size(); i++)
    {
        if(CastA[i] != 0)
        {
            Count = -1;
            break;
        }
    }
    
    return Count;
}