最佳节点 [华为]【树】

题目描述：

给定一个二叉树，每个节点有一个编号和一个值，该值可能为负数，请你找出一个最优节点（除根节点外），使得在该节点将树分成两棵树后（原来的树移除这个节点及其子节点，新的树以该节点为根节点），分成的两棵树各节点的和之间的差绝对值最大。请输出该节点编号，如有多个相同的差，输出编号最小的节点。
输入
4
4 9 -7 -8
0 1
0 3
1 2
第一行，四个节点，编号0-3，范围1-10000
第二行，节点0-3的权值
第三行到第五行，表示二叉树各节点之间的父子关系
0 1 //节点0的左节点是1
0 3 //节点0的右节点是3
1 2 //节点1的左节点是2
注意：左节点永远出现在右节点之前
输出
节点编号，示例中编号为3的节点是最优节点
样例1
输入
4
4 9 -7 -8
0 1
0 3
1 2
输出
3

思路：

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>  
using namespace std;
const int N=10010, INF=0x3f3f3f3f;

int n,sum;
int h[N],e[N],ne[N],w[N],idx;
int p[N];

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

int dfs(int t)
{
    int sum=0;
    for(int i=h[t];~i;i=ne[i])
    {
        int j=e[i];
        sum+=w[j];
        sum+=dfs(j);
    }
    return sum;
}

int main()
{
    scanf("%d",&n);
    for(int i=0;i<n;i++) p[i]=i;
    memset(h,-1,sizeof h);
    for(int i=0;i<n;i++)
    {
        scanf("%d",&w[i]);
        sum+=w[i];
    }
    for(int i=0;i<n-1;i++)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        add(a,b);
        p[b]=a;
    }
    int res=-INF;
    int point=0;
    for(int i=0;i<n;i++)
    {
        if(p[i]==i) continue;
        int t=dfs(i)+w[i];
        if(abs(sum-t-t)>res)
        {
            point=i;
            res=abs(sum-t-t);
        }
    }
    printf("%d\n",point,res);
    return 0;
}