[프로그래머스/Level 3/C++/월간 코드 챌린지 시즌2] 모두 0으로 만들기

https://school.programmers.co.kr/learn/courses/30/lessons/76503#

- 출처 : 프로그래머스 코딩 테스트 연습

프로그래머스

 

[문제 요약]

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

• 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다.

- 하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

- 트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다.

- 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 구하시오.

 

[문제 조건]

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

 

[문제 풀이]

 문제를 잘 생각해보면 아주 쉽게 풀리는 문제였다.

 

주어진 그래프가 트리의 형태일 때 잎이 아닌 노드를 0으로 만들려면 어떻게 해야되는지 한번 생각해보자.

 

잎이 아닌 노드를 0으로 만들기 위해서는 그 노드와 연결된 모든 노드들의 가중치의 합과 그 노드의 가중치 합이 0이 되야지만 모든 노드들을 조건에 따라 0으로 만들 수 있다.

 

그림으로 나타내면 다음과 같다.

 

현재 위치가 루트노드(0번)일 때 자식 노드들을 DFS 방식으로 탐색을 한다.

 

먼저 1번노드를 방문하면,  방문한 노드가 잎(Leaf)일경우 가중치를 부모노드에 리턴하고 부모노드에서 sum이라는 변수에 더해준다. 

 

다음으로 2번노드를 방문하면 , 방문한 노드가 잎이 아닐 경우 더 깊이 진행한다.

 

다음으로 3번 노드를 방문하면, 잎이기 때문에 2를 리턴하면 2번 노드에서의 sum에 누적시켜 준다. 여기서 sum의 초기값은 그 노드의 가중치로 초기화 되어 있다. 따라서 3번 노드를 탐색을 마친 상태면 2번 노드의 sum에는 3이란 값이 저장되어 있을 것이다.

 

이와 같이 4번 노드도 탐색하고 오면 2번 노드의 sum에는 최종적으로 5가 저장이 되어있다.

 

 

다시 돌아와 0번 노드에는 초기의 sum은 -5, 1번 노드를 탐색 후에는 -5, 2번 노드를 탐색을 마치면 -5+5해서 0이 저장된다. 

 

여기서 루트노드가 최종적으로 sum이 0이 되면 모든 노드들을 0으로 만들 수 있다는 뜻이고, 아니면 만들 수 없다는 뜻이다.

 

실제로 배열의 값이 바뀌는 것은 아니고 쉽게 이해하도록 하기위해 값을 변경한 것 뿐이니 오해하지 말자.

 

그 다음으로 연산 횟수는 각 노드에서 리턴해주는 sum의 값을 누적시켜 더 하면 그 값이 만들때 필요한 연산 횟수다.

 

다시 그림을 처음부터 확인해보면 3번,4번 노드 각각 0을 만들기 위해서 2번씩 총 4번

2번 노드는 3,4번 노드를 바꾸면서 같이 0이 됐을 것이니 넘기고

0번 노드에서 2번 노드를 0으로 바꾸면 5번의 연산이 필요하다 총 9번

0번 노드도 2번 노드를 바꾸면서 같이 0으로 바뀌었을 것이다.

 

따라서 총 9번의 연산으로 모두 0으로 만들 수 있다는 것이다.

 

[소스 코드]

#include <string>
#include <vector>

using namespace std;

bool flag = true;
long long cnt = 0;

long long dfs(int cur, int pnode, vector<vector<int>> &graph, vector<int> &a)
{
    long long sum = a[cur];
    
    for(int i=0; i<graph[cur].size(); i++)
    {
        int child = graph[cur][i];
        if(child == pnode) continue;

        sum += dfs(child, cur, graph, a);
    }

    cnt += abs(sum); 
    return sum;
}

long long solution(vector<int> a, vector<vector<int>> edges) {
    vector<vector<int>> graph(a.size());
    for(int i=0; i<edges.size(); i++)
    {
        graph[edges[i][0]].push_back(edges[i][1]);
        graph[edges[i][1]].push_back(edges[i][0]);
    }
    
    long long ans = dfs(0, -1, graph, a);
    
    if(ans == 0) return cnt;
    return -1;
}

참고로 주어지는 그래프의 형태를 띄고 있기때문에 역으로 올라가는 부분을 막아주면 별도의 visit배열을 만들지 않아도 된다.

굳