【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[N], p[N], s[N];
LL root, sum, ave;
vector<int>a[N];

LL avenum[N];
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)这两种情况就好了

*/