Codeforces ~ Garland ~ 767C （树形DP）

题意

给你一个N个结点的树，每个结点有一个权值。现在要将数分为三部分，使得三部分的权值相等，输出分割的点。
先输入N，然后接下来N行，每行输入 u 和 w，表示第 i 个结点和 u 有边，w表示第 i 个结点的权值。

思路

首先综合$sum\%3!=0$肯定不行。
dp[i]表示以 i 为根的树的权值和，dfs过程中统计子树权值即可，如果等于sum/3，那么就要从这个点上面分割开。
注意sum==0的时候，可能会只能分出来两部分。

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1e6+5;
int n, rt, sum, w[MAXN], dp[MAXN];
vector<int> G[MAXN];
vector<int> ans;
int dfs(int u, int fa)
{
    dp[u] = w[u];
    for (auto v: G[u])
    {
        if (v != fa)
            dp[u] += dfs(v, u);
    }
    if (dp[u] == sum/3) dp[u] = 0, ans.push_back(u);
    return dp[u];
}
int main()
{
    scanf("%d", &n);
    for (int v = 1; v <= n; v++)
    {
        int u; scanf("%d%d", &u, &w[v]);
        if (u)
        {
            G[u].push_back(v);
            G[v].push_back(u);
        }
        else rt = v;
        sum += w[v];
    }
    if (sum%3) { printf("-1\n"); return 0; }
    dfs(rt, -1);
    if (ans.size() <= 2) { printf("-1\n"); return 0; }
    printf("%d %d\n", ans[0], ans[1]);
    return 0;
}
/*
6
2 4
0 5
4 2
2 1
1 1
4 2
*/