【树形dp】CodeForces - 767C Garland

Problem Description

输入n代表有n个灯泡。接下来n行，第i行 给你i的父亲灯节点，和i这个灯的温度。
问你能不能用剪刀，剪两刀，将其分成三个部分，使得每个部分的温度和相等。
如果可以输出剪的边的儿子点。
否则输出-1。

思路：

自己强行用dfs模拟出来了，看了学长的代码突然发现自己弄麻烦了，因为它就是一棵树。
找出根节点，然后dfs遍历所有的点。在回溯的过程中，记录该子树的温度和。如果温度 是总和的 三分之一，那么记录该点，顺便让该点的温度变为0。这样就不影响寻找下一个点了。

#include<bits/stdc++.h>
using namespace std;
#define maxn 1000066
struct node
{
    int to, next;
};
node Map[maxn];
int head[maxn], cnt, sum[maxn], vis[maxn], total;
int ans[maxn], top;
void dfs(int u)
{
    vis[u] = 1;
    for(int i = head[u]; ~i; i = Map[i].next)
    {
        int to = Map[i].to;
        if(!vis[to])
        {
            dfs(to);
            sum[u] += sum[to];
        }
    }
    if(sum[u]*3 == total)//对于每个点，判断是否符合要求。
    {
        ans[top++] = u;
        sum[u] = 0;//为了不影响后面的点，变为0
    }
}
void add(int u, int v)
{
    Map[cnt] = (node)
    {
        v, head[u]
    };
    head[u] = cnt++;
}
int main()
{
    int n, u, i, num;
    while(~scanf("%d", &n))
    {
        memset(head, -1, sizeof(head));
        total = 0;
        cnt = 0;
        for(i = 1; i <= n; i++)
        {
            scanf("%d %d", &num, &sum[i]);
            if(num == 0) u = i;//找出根结点
            else add(num, i);//建父亲到儿子的边 即可
            total += sum[i];//记录总的温度
        }
        top = 0;
        memset(vis, 0, sizeof(vis));
        dfs(u);
        if(top < 2 || ans[1] == u)
            printf("-1\n");
        else printf("%d %d\n", ans[0], ans[1]);
    }
    return 0;
}