Codeforces Round #398 (Div. 2) C. Garland

http://codeforces.com/contest/767/problem/C

이 문제는 트리가 주어지고 트리의 각각의 노드에 온도가 주어질 때 노드들의 온도의 합이 같은 세 개의 덩어리로 나누는 문제이다. 나눌수 없다면 -1을 출력 있다면 나누는 두 지점을 (답이 여러개일 경우 아무거나) 출력 하는것이다.

이 문제는 크게 세 단계로 나눌 수 있다.

1.

루트 노드에서 dfs를 돌리면서 

w[i] = i번 노드를 포함하는 서브트리 의 모든 노드들의 온도의 합

을 구한다.

2.

w[v] 가 w[root] 의 3분의 1이 되는 지점 v들을 모두 구하여 조상들에게 모두 전달해 준다.

만약 노드 u 의 자식이 a,b,c 라고 했을때 a,b,c 를 각각 루트로 가지는 서브트리들에 대해 이러한 지점 v 를 가지는 자식들이 여러개라면 모두 갖고있는다(벡터 등을 이용), 또 a를 루트로 갖는 서브트리에 지점 v 가 여러개라면 하나만 갖고있는다.

3. 다시 dfs를 돌리면서 3분의 1지점인 v 가 2개이상이면 그 두점이 답이므로 출력후 종료, w[p] 가 w[root] 의 3분의 2가 되는 지점 p를 구한뒤 p의 자손중에 2번 에서 구한 v 가 하나라도있으면 그것과 p가 답.

단, 한가지 처리해줘야 할것이 root 노드는 자를 수 없기 때문에 root인 경우는 예외처리를 해줘야한다.

#include <vector>
#include <cstdio>
#define REP(i,a,b) for(int i=a;i<=b;++i)
#define FOR(i,n) for(int i=0;i<n;++i)
#define pb push_back
#define all(v) (v).begin(),(v).end()
#define sz(v) ((int)(v).size())
#define inp1(a) scanf("%d",&a)
#define inp2(a,b) scanf("%d%d",&a,&b)
#define inp3(a,b,c) scanf("%d%d%d",&a,&b,&c)
#define inp4(a,b,c,d) scanf("%d%d%d%d",&a,&b,&c,&d)
#define inp5(a,b,c,d,e) scanf("%d%d%d%d%d",&a,&b,&c,&d,&e)
using namespace std;
typedef vector<int> vi;
const int MAX_N = 1000002;
 
int N,a,b,c,root;
vi G[MAX_N];
int w[MAX_N];
vi third[MAX_N];
void dfs(int here){
    for(int there : G[here]){
        dfs(there);
        w[here]+=w[there];
    }
}
int dfs2(int here){
    for(int there : G[here]){
        int ret = dfs2(there);
        if(ret) third[here].pb(ret);
    }
    if(sz(third[here])>=2){
        printf("%d %d",third[here][0],third[here][1]);
        exit(0);
    }
    return (here!=root && w[root]==w[here]*3)?here:(sz(third[here])?third[here][0]:0);
}
 
void dfs3(int here){
    if(here!=root && w[here]==w[root]*2/3 && sz(third[here])){
        printf("%d %d",here,third[here][0]);
        exit(0);
    }
    for(int there : G[here]){
        dfs3(there);
    }
}
int main() {
    inp1(N);    
    REP(i,1,N){
        inp2(a,b);
        if(a) G[a].pb(i);
        else root=i;
        w[i]=b;
    }
    dfs(root);
    if(w[root]%3){
        printf("-1");
        return 0;
    }
    dfs2(root);
    dfs3(root);
    printf("-1");
    return 0;
}