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;
}