hihocoder - 1479 树形dp

题意：

给出一个有根树，每个节点都有一个权值，在其中选出不是根节点的两个节点，将这两个节点与其父亲之间边删去，可以将这棵树分成三个部分，求一共有多少种分法，可以让这三个部分的权值和相等。

思路：

这题一开始看错了题怎么都过不了，题目中说父亲节点为0的就是根节点，而我看成了0就是根节点，最后看别人代码才发现自己题读错了，太蠢。
树形dp，很显然如果存在这样的分法，所有节点的权值总和sumval%3==0。另外我们可以发现，那么可以用sum[u]表示以u为根的子树的权值和，这样dp一遍就能求出sum，如果存在多个sum[u]==sumval/3，那么我们从中任意选择两个不是根的节点不久可以分成三份了么？
然而这题目有坑，不会这么简单，如果可以这样考虑的前提是每个节点的权值一定大于0，但是这里权值可以非正，所以有可能存在父亲和儿子sum值一样大的情况，而同时选择这两个的方案显然就不对了。这里需要再进行一次dp。
这里dp[u]表示的是以节点u为根的子树一共有多少个sum[v]==sumval/3的节点。首先要想到，无论权值正负，只要是在同一父亲的兄弟之间选择子树是互不影响的，这样在兄弟v1和v2之间选择一共就用dp[v1]*dp[v2]种方案。另外，选择一个父亲u和一个后代节点v也是有可能的，只要这个父亲的权值和sum[u]==sumval/3*2，这样可以满足条件，这样的方案数有dp[v]个，基于这样的思想，一遍dfs即可求出答案。

代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = 1e6 + 10;

vector <int> tree[MAXN];
ll val[MAXN], sum[MAXN], dp[MAXN];
ll ans, sumval, root;

int dfs1(int u) {
sum[u] = val[u];
dp[u] = 0;
int cnt = tree[u].size();
for (int i = 0; i < cnt; i++) {
int v = tree[u][i];
sum[u] += dfs1(v);
dp[u] += dp[v];
}
if (sum[u] == sumval / 3) ++dp[u];
return sum[u];
}

void dfs2(int u) {
int cnt = tree[u].size();
ll tmp = 0;
for (int i = 0; i < cnt; i++) {
int v = tree[u][i];
ans += dp[v] * tmp;
tmp += dp[v];
if (u != root && sum[u] == sumval / 3 * 2)
ans += dp[v];
dfs2(v);
}
}

int main() {
//freopen("in.txt", "r", stdin);
int T;
scanf("%d", &T);
while (T--) {
int n;
scanf("%d", &n);
for (int i = 0; i <= n; i++) tree[i].clear();
sumval = 0;
for (int i = 1; i <= n; i++) {
int p;
scanf("%lld%d", &val[i], &p);
if (p == 0) root = i;
else tree[p].push_back(i);
sumval += val[i];
}
if (sumval % 3 != 0) {
puts("0");
continue;
}
ans = 0;
dfs1(root);
dfs2(root);
printf("%lld\n", ans);
}
return 0;
}