三等分
描述
小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
是个好题,codeforce做过一个类似的,不过只要求一个方案
#include <iostream> #include <cstdio> #include <cmath> #include <cstring> #include <algorithm> #include <vector> using namespace std; typedef long long LL; const int N=1e6+10,mod=20090717,inf=2e9+10; int T,n,v[N],root,f,all,dp[N],si,first,second; LL ans; vector<int > G[N]; void dfs(int u,int fa) { dp[u] = v[u]; for(int i = 0; i < G[u].size(); ++i) { int to = G[u][i]; if(to == fa) continue; dfs(to,u); dp[u] += dp[to]; } } void dfs(int u) { if(dp[u] == all) { ans += first + second; } if(dp[u] == all*2 && u != root) second++; for(int i = 0; i < G[u].size(); ++i) { int to = G[u][i]; dfs(to); } if(dp[u] == all) first++; if(dp[u] == all*2 && u != root) second--; } int main() { scanf("%d",&T); while(T--) { scanf("%d",&n); for(int i = 0; i <= n; ++i) G[i].clear(); all = 0; for(int i = 1; i <= n; ++i) { scanf("%d%d",&v[i],&f); if(f == 0) root = i; else G[f].push_back(i); } for(int i = 1; i <= n; ++i) all += v[i]; if((all % 3+3)%3 != 0) { puts("0"); continue; } all/=3;si = 0;ans = 0; first = second = 0; dfs(root,0); dfs(root); printf("%lld\n",ans); } return 0; }