题目大意
院子里有一颗又高又大的草莓树,草莓树有n个节点,每个节点都结了一个草莓,吃掉第i个结点的草莓可以得到ai的营养值,由于草莓可能会坏掉,所以ai可能为负值,也可能为0。
现在要砍掉这颗树的两条边,使树变成三份,并且使得三份各自草莓营养值的和恰好一样。请问是否有这样的方法呢?如果有,请输出YES,否则输出NO。
解题思路
树形dp。
首先如果所有值的和不是3的倍数就是NO,并且每一份的值必须要是所有草莓值和的1/3,记k为这个值。
所以统计一个子树的权值和,以及一个dp[i],表示i为根的子树中是否存在一个权值为k的子树。
有解的情况分为以下两种:
1. 如果当前结点i的儿子中,有两个儿子s1,s2,dp[s1]和dp[s2]都被标记为true。此时切掉s1和s2的父边为一个可行方案。
2. 如果当前结点不是根,子树和恰好为2*k,并且其子树中存在权值为k的子树。此时切掉其父边,以及它某个后代的父边,形成一个可行的方案。
只要判断是否存在这两种情况即可 [引用]
#include <cstdio>
#include <iostream>
#include <cstring>
#include <string>
#include <algorithm>
using namespace std;
const int N = 1000005;
int n, w[N], dp[N], head[N];
long long first, second, root, val, idc;
long long ans;
struct ln{
int to, next;
} ed[N*2];
void adde(int u, int v){
ed[++idc].to = v;
ed[idc].next = head[u];
head[u] = idc;
}
void dfs(int u,int fa) {
dp[u] = w[u];
for(int i = head[u]; i; i = ed[i].next) {
int v = ed[i].to;
if(v != fa) {
dfs(v, u);
dp[u] += dp[v];
}
}
}
void dfs(int u) {
if(dp[u] == val) {
ans += first + second;
}
if(dp[u] == (val << 1) && u != root) second++;
for(int i = head[u]; i; i = ed[i].next) {
int v = ed[i].to;
dfs(v);
}
if(dp[u] == val) first++;
if(dp[u] == (val << 1) && (u != root)) second--;
}
int main() {
int T; scanf("%d", &T);
while( T-- ){
scanf("%d",&n);
memset(head, 0, sizeof (int) * (n + 1));
memset(ed, 0, sizeof (int) * (n + 1));
memset(dp, 0, sizeof (int) * (n + 1));
memset(w, 0, sizeof (int) * (n + 1));
val = 0; idc = 0;
for(int i = 1; i <= n; i++) {
int ff;
scanf("%d%d", &ff, &w[i]);
val += w[i];
if( !ff ) root = i;
else adde(ff, i);
}
if( ((val % 3) + 3) % 3 ) { puts("NO"); continue; } //val可能为负
val /= 3; ans = 0;
first = second = 0;
dfs(root,0);
dfs(root);
ans ? puts("YES") : puts("NO");
}
}