strawberry 分草莓 （树上DP）

题目大意
院子里有一颗又高又大的草莓树，草莓树有n个节点，每个节点都结了一个草莓，吃掉第i个结点的草莓可以得到ai的营养值，由于草莓可能会坏掉，所以ai可能为负值，也可能为0。
现在要砍掉这颗树的两条边，使树变成三份，并且使得三份各自草莓营养值的和恰好一样。请问是否有这样的方法呢？如果有，请输出YES，否则输出NO。

解题思路
树形dp。
首先如果所有值的和不是3的倍数就是NO，并且每一份的值必须要是所有草莓值和的1/3，记k为这个值。
所以统计一个子树的权值和，以及一个dp[i]，表示i为根的子树中是否存在一个权值为k的子树。
有解的情况分为以下两种：
1. 如果当前结点i的儿子中，有两个儿子s1,s2，dp[s1]和dp[s2]都被标记为true。此时切掉s1和s2的父边为一个可行方案。
2. 如果当前结点不是根，子树和恰好为2*k，并且其子树中存在权值为k的子树。此时切掉其父边，以及它某个后代的父边，形成一个可行的方案。
只要判断是否存在这两种情况即可 [引用]

#include <cstdio>
#include <iostream>
#include <cstring>
#include <string>
#include <algorithm>
using namespace std;

const int N = 1000005;

int n, w[N], dp[N], head[N];
long long first, second, root, val, idc;
long long ans;

struct ln{
    int to, next;
} ed[N*2];

void adde(int u, int v){ 
    ed[++idc].to = v;
    ed[idc].next = head[u];
    head[u] = idc; 
}

void dfs(int u,int fa) {
    dp[u] = w[u];
    for(int i = head[u]; i; i = ed[i].next) {
        int v = ed[i].to;
        if(v != fa) {
            dfs(v, u);
            dp[u] += dp[v];
        }
    }
}

void dfs(int u) {
    if(dp[u] == val) {
        ans += first + second;
    }
    if(dp[u] == (val << 1) && u != root) second++;
    for(int i = head[u]; i; i = ed[i].next) {
        int v = ed[i].to;
        dfs(v);
    }
    if(dp[u] == val) first++;
    if(dp[u] == (val << 1) && (u != root)) second--;
}

int main() {
    int T; scanf("%d", &T);
    while( T-- ){
        scanf("%d",&n);
        memset(head, 0, sizeof (int) * (n + 1));
        memset(ed, 0, sizeof (int) * (n + 1));
        memset(dp, 0, sizeof (int) * (n + 1));
        memset(w, 0, sizeof (int) * (n + 1));
        val = 0; idc = 0;
        for(int i = 1; i <= n; i++) {
            int ff;
            scanf("%d%d", &ff, &w[i]);
            val += w[i];
            if( !ff ) root = i;
            else adde(ff, i);
        }    
        if( ((val % 3) + 3) % 3 ) {  puts("NO"); continue; }  //val可能为负 
        val /= 3; ans = 0;
        first = second = 0;
        dfs(root,0);
        dfs(root);
        ans ? puts("YES") : puts("NO");
    }
}