[BOJ 26972] Barn Tree

https://www.acmicpc.net/problem/26972

26972번: Barn Tree

 

- 문제 요약

 

Farmer John의 농장에는 건초가 들어있는 N개의 헛간이 있습니다. N개의 헛간에는 1~N까지 번호가 붙어있습니다.

또한, 두 개의 헛간을 이어주는 길이 N-1개 있습니다.

John은 N개의 헛간에 들어있는 건초의 수를 동일하게 만들고 싶어 합니다.

헛간과 헛간을 이어주는 길을 통해 건초를 옮길 수 있습니다.

현재 헛간에 있는 건초 수보다 작거나 같은 수의 건초만 옮길 수 있습니다.

가능한 최소 주문 수로 작업을 완료하기 위해 Farmer John이 말할 수 있는 주문을 각 줄에 출력하세요.

 

 

- 알고리즘 정리

 

dfs 활용으로 풀면 되는 문제입니다. up으로 그래프 순회 한 번 하고, down으로 순회 한 번 하면서 과정 중에 얻는 주문을 ans 벡터에 넣어줍니다. (ans 벡터의 사이즈와 값을 출력하면 풀이 완료)

 

 

- 코드 작성

 

#include<bits/stdc++.h>
#include<tuple>
using namespace std;

#define MAX 200005
typedef long long ll;
int n,di[4]={0,-1,0,1},dj[4]={-1,0,1,0};
ll avg,h[MAX],sum[MAX],sz[MAX];
vector<int>v[MAX];
vector<tuple<int,int,ll>>ans;

void dfs_up(int x,int p){
    sz[x]=1;
    sum[x]=h[x];
    for(auto &i:v[x]){
		if(i!=p){
	        dfs_up(i,x);
	        sum[x]+=sum[i];
	        sz[x]+=sz[i];
	    }
    }
    if(x!=1&&sum[x]>avg*sz[x]){
        ans.push_back({x,p,sum[x]-avg*sz[x]});
        sum[p]+=sum[x]-avg*sz[x];
        sum[x]-=sum[x]-avg*sz[x];
    }
}

void dfs_down(int x,int p){
    for(auto &i:v[x]){
		if(i!=p){
	        if(sum[i]<avg*sz[i]){
	            ans.push_back({x,i,avg*sz[i]-sum[i]});
	            sum[x]-=avg*sz[i]-sum[i];
	            sum[i]+=avg*sz[i]-sum[i];
	        }
	        dfs_down(i,x);
    	}
    }
}

int main(){
    ios::sync_with_stdio(false);
	cin.tie(NULL);
    cin>>n;
    for(int i=1;i<=n;i++){
        cin>>h[i];
        avg+=h[i];
    }
    avg/=n;
    for(int i=1;i<n;i++){
        int p,q;
		cin>>p>>q;
        v[p].push_back(q);
        v[q].push_back(p);
    }
    dfs_up(1,0);
    dfs_down(1,0);
    cout<<ans.size()<<'\n';
    for(auto &[i,j,k]:ans){
    	cout<<i<<' '<<j<<' '<<k<<'\n';
	}
}

코드 제출 결과

코드가 너무 보기에 난잡해서 수정 한 번 했는데 컴파일 에러에 타임 아웃에.....

(뭘 잘못한 지도 모르고 뭐가 문제지 하면서 맞왜틀 계속 시전)

알고 보니 endl이랑 '\n' 시간 차이 때문에 틀렸던 것......

 

+) #define endl '\n'을 생활화합시다