【[Offer收割]编程练习赛9 C】三等分

【题目链接】:http://hihocoder.com/problemset/problem/1479

【题意】

.

【题解】

首先算出所有节点的权值的和val;
然后如果val%3!=0则直接输出0；
否则令key = val/3
对于当前节点x;
有两种情况;
一种是
两个要删掉的节点都在同一棵子树里面;
则某一个子树里面必须有一个子树和为key的节点;同时这个点x的子树和为2*key;->注意这时x节点不能为根节点
另外一种情况就是;
两个key的节点分别在x的两棵子树里面;

【Number Of WA】

1(忘记初始化vector了)

【完整代码】

#include <bits/stdc++.h>
using namespace std;
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define LL long long
#define rep1(i,a,b) for (int i = a;i <= b;i++)
#define rep2(i,a,b) for (int i = a;i >= b;i--)
#define mp make_pair
#define pb push_back
#define fi first
#define se second
#define ms(x,y) memset(x,y,sizeof x)

typedef pair<int,int> pii;
typedef pair<LL,LL> pll;

const int dx[9] = {0,1,-1,0,0,-1,-1,1,1};
const int dy[9] = {0,0,0,-1,1,-1,1,-1,1};
const double pi = acos(-1.0);
const int N = 1e5+100;

int n,root,a[N],fa,key;
vector <int> son[N];
LL dp1[N],sum[N],ans;

void dfs(int x)
{
    dp1[x] = 0;sum[x] = a[x];
    for (int y:son[x])
    {
        dfs(y);
        sum[x]+=sum[y];
        ans += dp1[x]*dp1[y];
        dp1[x]+=dp1[y];
    }
    if (x!=root && sum[x]==key*2) ans+=dp1[x];
    if (sum[x]==key) dp1[x]++;
}

int main()
{
    //freopen("F:\\rush.txt","r",stdin);
    ios::sync_with_stdio(false),cin.tie(0);//scanf,puts,printf not use
    int T;
    cin >> T;
    while (T--)
    {
        ans = 0,key = 0;
        cin >> n;
        rep1(i,1,n) son[i].clear();
        rep1(i,1,n)
        {
            cin >> a[i] >> fa;
            if (fa==0)
                root = i;
            else
                son[fa].pb(i);
            key+=a[i];
        }
        if (key%3!=0)
        {
            cout << 0 << endl;
            continue;
        }
        key/=3;
        dfs(root);
        cout << ans << endl;
    }
    return 0;
}