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

문제

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

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

 

풀이

1. +, -하는 가중치가 같기에 모든 가중치의 합이 0이여야 결국 0으로 만들 수 있습니다. 따라서 0이 아니면 -1을 반환합니다.

 

2. 조금만 생각해본다면 연결된 두 점에 대해서 연산할 때, 리프부터 처리해서 루트로 올라와야 빠짐없이 0으로 만들 수 있습니다.

 

3. 트리임이 보장되었기에 0번노드를 루트로 잡고 dfs를 돌리는데 리프노드에선 연산을 하지 않도록 처리합니다. 이는 노드들을 양 방향 간선으로 이어주었기에 리프 노드의 간선 개수가 1개임을 생각하면 간단히 처리할 수 있습니다.

 

코드

#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
ll ans, s;
vector<ll> a;
vector<int> adj[300001];
 
void go(int u, int p) {
	for (int v : adj[u]) {
		if (v == p) continue;
		if (adj[v].size() != 1) go(v, u);
		ans += abs(a[v]);
		a[u] += a[v], a[v] = 0;
	}
}
 
ll solution(vector<int> A, vector<vector<int>> edges) {
	for (int i : A) a.push_back((ll)i);
	for (auto i : edges) {
		adj[i[0]].push_back(i[1]);
		adj[i[1]].push_back(i[0]);
	}
	bool allz = 1;
	for (ll i : a) s += i, allz &= (i == 0);
	if (allz) return 0;
	if (s != 0) return -1;
	go(0, -1);
	return ans;
}