Hihocoder #1479 : 三等分 树形DP

http://hihocoder.com/problemset/problem/1479

描述

小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，表示方案的数量。

----------------------------------------------------------------

显然是个树DP，首先我们dfs求个子树权值和NUM

其次，我们要找的肯定是值为sum的1/3的点，即1/3点。记为第一类点first

值为2/3的点记为 second

则代码为：

void dfs2(int x,int fa)
{
    if (num[x]==all)//遇到三分点，统计答案
        ans+=first+second;			//注意遇到1/3点，计算完答案后不可马上return，原因是本题有负的数值，因此还需要往下跑继续找1/3点
    if (num[x]==all*2&&x!=root) second++;//统计2/3点
    for(int i=0;i<mp[x].size();i++)
    {
        int v=mp[x][i];
        if (x==fa) continue;
        dfs2(v,x);
    }
    if (num[x]==all) first++;//累计1/3点
    if (num[x]==all*2&&x!=root) second--;//（显然，如果不是当前点的祖先的2/3点是不会被记录的,准确地说,是会被去掉）
}

当我们递归时， 遇到某个点是1/3点，则 ans+=1/3点数量+2/3的数量

前者好理解，别的1/3点会和当前点组成三等分点对，显然符合条件。

那么2/3怎么理解呢？是这样的，由于递归的原因，如果当前点之前存在2/3点，则该点一定是当前点的非根祖先。（显然，如果不是当前点的祖先的2/3点是不会被记录的） 

 因此那个非根祖先和当前点 也能把树三等分。

就两种情况，

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int v[100000+50];
vector<int >mp[100000+50];
int num[100000+50];

void dfs(int x,int fa)
{
    num[x]=v[x];
    for(int i=0; i<mp[x].size(); i++)
    {
        int v=mp[x][i];
        if (v==fa)continue;
        dfs(v,x);
        num[x]+=num[v];
    }
}
ll first,second,ans,all,root;
void dfs2(int x,int fa)
{
    if (num[x]==all)
        ans+=first+second;
    if (num[x]==all*2&&x!=root) second++;
    for(int i=0; i<mp[x].size(); i++)
    {
        int v=mp[x][i];
        if (x==fa) continue;
        dfs2(v,x);
    }
    if (num[x]==all) first++;
    if (num[x]==all*2&&x!=root) second--;
}
int main()
{

    int t;
    cin>>t;
    while(t--)

    {
        memset(num,0,sizeof num);
        int n;
        cin>>n;
        for(int i=1; i<=n; i++)mp[i].clear();
        int vv,f;
        all=ans=first=second=0;
        for(int i=1; i<=n; i++)
        {
            scanf("%d%d",&v[i],&f);
            all+=v[i];
            if (f==0)
                root=i;
            else
                mp[f].push_back(i);
        }
        if (all%3!=0)
        {
            printf("0\n");
            continue;
        }
        all/=3;

        dfs(root,0);
        dfs2(root,0);

        printf("%lld\n",ans);
    }

}