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

题目3 : 三等分

时间限制:10000ms
单点时限:1000ms
内存限制:256MB

描述

小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

#include<stdio.h>
#include<iostream>
#include<string.h>
#include<string>
#include<ctype.h>
#include<math.h>
#include<set>
#include<map>
#include<vector>
#include<queue>
#include<bitset>
#include<algorithm>
#include<time.h>
using namespace std;
void fre() { freopen("c://test//input.in", "r", stdin); freopen("c://test//output.out", "w", stdout); }
#define MS(x, y) memset(x, y, sizeof(x))
#define ls o<<1
#define rs o<<1|1
typedef long long LL;
typedef unsigned long long UL;
typedef unsigned int UI;
template <class T1, class T2>inline void gmax(T1 &a, T2 b) { if (b > a)a = b; }
template <class T1, class T2>inline void gmin(T1 &a, T2 b) { if (b < a)a = b; }
const int N = 1e6 + 10, M = 0, Z = 1e9 + 7, inf = 0x3f3f3f3f;
template <class T1, class T2>inline void gadd(T1 &a, T2 b) { a = (a + b) % Z; }
int casenum, casei;
int n;
int v
, p
, s
;
LL root, sum, ave;
vector<int>a
;

LL avenum
;
LL ans;

void dfs(int x)
{
s[x] = v[x];
avenum[x] = 0;
for (auto y : a[x])
{
dfs(y);
s[x] += s[y];
avenum[x] += avenum[y];
}
if (x == root)return;
if (s[x] == ave * 2)
{
ans += avenum[x];
}
if (s[x] == ave)
{
ans -= avenum[x];
++avenum[x];
}
}
LL solve()
{
scanf("%d", &n);
root = 0;
sum = 0;
for (int i = 0; i <= n; ++i)
{
a[i].clear();
avenum[i] = s[i] = 0;
}
for (int i = 1; i <= n; ++i)
{
scanf("%d%d", &v[i], &p[i]);
if (p[i] == 0)root = i;
else a[p[i]].push_back(i);
sum += v[i];
}
if (sum % 3)return 0;
ave = sum / 3;

ans = 0;
dfs(root);
ans += avenum[root] * (avenum[root] - 1) / 2;

return ans;
}
int main()
{
scanf("%d", &casenum);
for (casei = 1; casei <= casenum; ++casei)
{
printf("%lld\n", solve());
}
return 0;
}
/*
【trick&&吐槽】
20分钟写好代码结果我以1为根了QAQ 结果最后时间才AC，不然就rank1 了 好气啊！

【分析】
只要考虑(1/3)外套 (1/3) 和(2/3) 内嵌(1/3)这两种情况就好了

*/