CodeForces 767C :Garland 树形DP

传送门

题目描述

n节点的树,每个节点都有一个权值ti。
现在让你恰好切掉其中的两条边,使得构成的新的3棵树的权值和都变为一样的。
即sum1=sum2=sum3.

分析

挺简单的一个树形DP，就是有个细节没太注意一直wa
我们去求每一个节点的子树的大小，如果大小等于sum/3就记录答案，然后子树大小置0
需要注意的是，我们只需要记录两个答案就可以了就比如我从根节点出发，已经删减了两棵子树，这样根节点形成的子树的也是sum/3，但这个时候根节点也有可能被记录，比如根节点的数值是0，这样他的子树大小也是满足条件的，但是如果记录这个节点，答案就不正确了

代码

#include <iostream>
#include <cstdio>
#include <cmath>
#include <algorithm>
#include <queue>
#include <cstring>
#define debug(x) cout<<#x<<":"<<x<<endl;
#define _CRT_SECURE_NO_WARNINGS
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> PII;
const int INF = 0x3f3f3f3f;
const int N = 1e6 + 10;
int a[N];
int dp[N];
int x1,x2;
int h[N],ne[N * 2],e[N * 2],idx;
int n;
int sum;

void add(int x,int y){
    ne[idx] = h[x],e[idx] = y,h[x] = idx++;
}

void dfs(int u,int fa){
    dp[u] = a[u];
    for(int i = h[u];~i;i = ne[i]){
        int j = e[i];
        if(j == fa) continue;
        dfs(j,u);
        dp[u] += dp[j];
        if(dp[j] == sum / 3){
            if(!x1) x1 = j;
            else (!x2) x2 = j;
            dp[u] -= dp[j];
        }
    }
}

int main(){
    memset(h,-1,sizeof h);
    scanf("%d",&n);
    int root;
    for(int i = 1;i <= n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        a[i] = y;
        if(!x) root = i;
        else add(x,i),add(i,x);
        sum += y;
    }
    if(sum % 3){
        puts("-1");
        return 0;
    }
    dfs(root,-1);
    if(!x1 || !x2){
        puts("-1");
        return 0;
    }
    printf("%d %d\n",x1,x2);
    return 0;
}

/**
*　　┏┓　　　┏┓+ +
*　┏┛┻━━━┛┻┓ + +
*　┃　　　　　　　┃
*　┃　　　━　　　┃ ++ + + +
*  ████━████+
*  ◥██◤　◥██◤ +
*　┃　　　┻　　　┃
*　┃　　　　　　　┃ + +
*　┗━┓　　　┏━┛
*　　　┃　　　┃ + + + +Code is far away from 　
*　　　┃　　　┃ + bug with the animal protecting
*　　　┃　 　 ┗━━━┓ 神兽保佑,代码无bug　
*　　　┃ 　　　　　　 ┣┓
*　　  ┃ 　　　　　 　┏┛
*　    ┗┓┓┏━┳┓┏┛ + + + +
*　　　　┃┫┫　┃┫┫
*　　　　┗┻┛　┗┻┛+ + + +
*/