Hihocoder #1479 : 三等分 树形DP

三等分 

描述

小Hi最近参加了一场比赛，这场比赛中小Hi被要求将一棵树拆成3份，使得每一份中所有节点的权值和相等。

比赛结束后，小Hi发现虽然大家得到的树几乎一模一样，但是每个人的方法都有所不同。于是小Hi希望知道，对于一棵给定的有根树，在选取其中2个非根节点并将它们与它们的父亲节点分开后，所形成的三棵子树的节点权值之和能够两两相等的方案有多少种。

两种方案被看做不同的方案，当且仅当形成方案的2个节点不完全相同。

输入

每个输入文件包含多组输入，在输入的第一行为一个整数T，表示数据的组数。

每组输入的第一行为一个整数N，表示给出的这棵树的节点数。

接下来N行，依次描述结点1~N，其中第i行为两个整数Vi和Pi，分别描述这个节点的权值和其父亲节点的编号。

父亲节点编号为0的节点为这棵树的根节点。

对于30%的数据，满足3<=N<=100

对于100%的数据，满足3<=N<=100000, |Vi|<=100, T<=10

输出

对于每组输入，输出一行Ans，表示方案的数量。

样例输入

2
3
1 0
1 1
1 2
4
1 0
1 1
1 2
1 3

样例输出

1
0

 

是个好题，codeforce做过一个类似的，不过只要求一个方案

#include <iostream>
#include <cstdio>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
typedef long long LL;
const int N=1e6+10,mod=20090717,inf=2e9+10;

int T,n,v[N],root,f,all,dp[N],si,first,second;
LL ans;
vector<int > G[N];
void dfs(int u,int fa) {
    dp[u] = v[u];
    for(int i = 0; i < G[u].size(); ++i) {
        int to = G[u][i];
        if(to == fa) continue;
        dfs(to,u);
        dp[u] += dp[to];
    }
}
void dfs(int u) {
    if(dp[u] == all) {
        ans += first + second;
    }
    if(dp[u] == all*2 && u != root) second++;
    for(int i = 0; i < G[u].size(); ++i) {
        int to = G[u][i];
        dfs(to);
    }
    if(dp[u] == all) first++;
    if(dp[u] == all*2 && u != root) second--;
}
int main() {
    scanf("%d",&T);
    while(T--) {
        scanf("%d",&n);
        for(int i = 0; i <= n; ++i) G[i].clear();
        all = 0;
        for(int i = 1; i <= n; ++i) {
            scanf("%d%d",&v[i],&f);
            if(f == 0) root = i;
            else G[f].push_back(i);
        }
        for(int i = 1; i <= n; ++i) all += v[i];
        if((all % 3+3)%3 != 0) {
            puts("0");
            continue;
        }
        all/=3;si = 0;ans = 0;
        first = second = 0;
        dfs(root,0);
        dfs(root);
        printf("%lld\n",ans);

    }
    return 0;
}