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

样例输入

231 01 11 241 01 11 21 3

样例输出

10

#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|1typedef 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)这两种情况就好了*/