본문 바로가기

Algorithm/DFS

(C++) - 프로그래머스(월간 코드 챌린지 시즌2) : 모두 0으로 만들기

반응형

programmers.co.kr/learn/courses/30/lessons/76503

 

코딩테스트 연습 - 모두 0으로 만들기

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다. 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한

programmers.co.kr

dfs로 푼 문제였습니다.

 

풀이방법

 1. cost를 모두 더했을 때 0이 아니라면 -1을 반환합니다.

 

 2. 그래프를 인접리스트로 재구성합니다.

 

 3. dfs를 시행합니다. 어떤 정점에서 시작해도 모든 정점을 확인하기 때문에 정답을 구할 수 있습니다. 연결된 모든 정점을 확인하면서 다음 정점의 cost 절댓값을 ans변수에 더해줍니다. 현재 정점의 cost는 next로부터 오는 모든 cost의 합입니다. 이에 맞춰 현재 정점 cost를 갱신해줍니다. 

 

 4. 답을 반환합니다.

 

Code

#include <bits/stdc++.h>
#define ll long long
using namespace std;
int ck[300001];
ll cost[300001],ans;
void dfs(int cur, vector<int>graph[]){
    ck[cur] = 1;
    for(int next : graph[cur]){
        if(ck[next]) continue;
        dfs(next,graph);
        ans += abs(cost[next]);
        cost[cur] += cost[next];
    }
}

ll solution(vector<int> a, vector<vector<int>> edges) {
    ll sum = 0;
    ans = 0;
    vector<int> graph[a.size()];
    memset(ck,0,sizeof(ck));
    memset(cost,0,sizeof(cost));
    for(int i = 0; i < a.size(); i++) {
        cost[i] = a[i], sum += cost[i];
    }
    if(sum) return -1;
    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]);
    }
    dfs(0, graph);
    return ans;
}