CoderForces 767 C （树形dp）

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/imzxww/article/details/81669876

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

题意：

一棵树，每个节点有一个权值t，把它分成三部分，每部分要求权值和相等。

思路：

直接dfs往下找，如果这里的权值之和等于目标的话，就把子树的权值设为0，不影响上面的。

代码:

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int maxn=1e6+10;
vector<int>e[maxn];
int sz[maxn];
int w[maxn];
int sum;
int ans[maxn];
int cnt;
void dfs(int x,int pre)
{
    sz[x]=w[x];
    for(int i=0;i<e[x].size();i++)
    {
        int v=e[x][i];
        if(v==pre)continue;
        dfs(v,x);
        sz[x]+=sz[v];
    }
    if(sz[x]==sum)ans[++cnt]=x,sz[x]=0;
}
int main()
{
    int n;
    scanf("%d",&n);
    int root;
    for(int i=1;i<=n;i++)
    {
        int x;
        scanf("%d%d",&x,&w[i]);
        if(x)e[x].push_back(i),e[i].push_back(x);
        else root=i;
        sum+=w[i];
    }
    if(sum%3!=0)
    {
        printf("-1\n");
    }
    else
    {
        sum/=3;
        dfs(root,-1);
        if(cnt<3)printf("-1\n");
        else printf("%d %d\n",ans[1],ans[2]);
    }
}