[机房练习赛4.7] 分草莓 树上DP

分草莓（strawberry.in/strawberry.out）

院子里有一颗又高又大的草莓树，草莓树有n个节点，每个节点都结了一个草莓，吃掉第i个结点的草莓可以得到ai的营养值，由于草莓可能会坏掉，所以ai可能为负值，也可能为0。
现在要砍掉这颗树的两条边，使树变成三份，并且使得三份各自草莓营养值的和恰好一样。请问是否有这样的方法呢？如果有，请输出YES，否则输出NO。
【输入格式】
第一行一个数t，表示测试点的个数。
接下来t组：
每组第一行一个数n，表示结点的个数。
接下来n行，每行两个数fai和ai，表示第i个结点的父亲是fai，第i个结点有营养值为ai的草莓，根节点的fai记为0.
【输出格式】
输出t行，如果第t组有解，则输出YES，否则输出NO
【输入样例】
2
6
2 4
0 5
4 2
2 1
1 1
4 2
6
2 4
0 6
4 2
2 1
1 1
4
4 2
【输出样例】
YES
NO
【样例解释】
第一组可以切掉1号和4号的父边，形成均等的三份。
【数据规模】
10% 数据满足3≤n≤100。
100% 数据满足t≤20，3≤n≤100000。-100≤ai≤100
【解法】
树形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的子树。此时切掉其父边，以及它某个后代的父边，形成一个可行的方案。
只要判断是否存在这两种情况即可。

/*由于数据太水，我这个大暴力过了，其实我这个方法与上面略有不同，他先在树上跑一次dfs，算出所有点及其子节点的营养值。枚举所有这样的点，分别求lca（我不知道这样还能过）*/
#include <cstdio>
#include <algorithm>
using namespace std;

const int N = 500010;
const int M = N + N;
const int P = 16;

inline int read() {
    int x = 0, f = 1;char ch = getchar();
    while( ch < '0' || ch > '9'){if(ch == '-')f=-1;ch=getchar();}
    while(ch>='0' && ch <='9'){x = x*10+ch-'0';ch = getchar();}
    return x*f;
}
int n, T;
int head[N], dest[M], last[M], dast[M], etot;
int dep[N], anc[N][P+1], sum[N], in[N];

void adde( int u, int v ) {
    dest[++etot] = v;last[etot] = head[u];head[u] = etot;
}
void dfs( int u, int f ) {
    anc[u][0] = f;
    for( int p = 1; p <= P; p++ )
        anc[u][p] = anc[anc[u][p-1]][p-1];
    for( int t = head[u]; t; t = last[t] ) {
        int v = dest[t];
        if( v == f ) continue;
        dep[v] = dep[u] + 1;
        dfs( v, u );
        sum[u] += sum[v];
    }
}
int lca( int u, int v ) {
    if( dep[u] < dep[v] ) swap(u,v);
    int t = dep[u] - dep[v];
    for( int p = 0; t; t>>=1,p++ )
        if( t & 1 ) u = anc[u][p];
    if( u == v ) return u;
    for( int p = P; p >= 0; p-- ) 
        if( anc[u][p] != anc[v][p] )
            u = anc[u][p], v = anc[v][p];
    return anc[u][0];
}
int main() {
    freopen("strawberry.in","r",stdin);
    freopen("strawberry.out","w",stdout);
    scanf("%d", &T);
    while( T-- ){
        memset(head,0,sizeof(head));
        memset(anc,0,sizeof(anc));
        memset(sum,0,sizeof(sum));
        n = read();
        for( register int i = 1; i <= n; i++ ){
            int x,y;
            x = read(); y = read();
            adde( i, x ); adde( x, i );
            sum[i] = y;in[i]++;
        } int ddf;
        dfs(0,0);
        int yu = sum[0]/3;
        int q[10000],tail=0;
        for( register int i = 1; i <= n; i++ ){
            if( sum[i] == yu ){
                tail++;
                q[tail] = i;
            }
        }
        int fg,jk = 0;
        for( register int i = 2; i <= tail; i++ )
            for( int j = 1; j < i; j++ )
                if( i != j && jk == 0 ){
                    int fg = lca( q[i], q[j] );
                    if( fg != q[i] && fg != q[j] ){
                        jk = 1;break;
                    }
                }
        if( jk ) {
            printf("YES\n");continue;
        } else{
            printf("NO\n"); continue;
        }
    }
    return 0;
}

std

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<string>
#include<cmath>
#include<queue>
#include<algorithm>
#include<iostream>
#define rep(i, a, b) for(int i = a; i <= b; i++)
#define repd(i, a, b) for(int i = b; i >= a; i--)
#define sz(x) (int)((x).size())
#define ms(x, y) memset(x, y, sizeof(x))
#define pb push_back
#define ll long long 
#define mp make_pair
using namespace std;


const int N = 100005;


int fa[N];
int dis[N];

int n;
int t[N];
int s[N];
int dp[N];

struct edge {
    int v, nxt;
}e[N];
int b[N];
int en;

void addedge(int u, int v) {
    e[en].v = v;
    e[en].nxt = b[u];
    b[u] = en++;
}
vector<int> xu;
queue<int> q;
void bfs(int rt) {
    while(!q.empty())q.pop();
    q.push(rt);
    xu.pb(rt);
    dis[rt] = 1;
    while(!q.empty()) {
        int now = q.front();
        q.pop();
        for(int i = b[now]; i!=-1; i = e[i].nxt){
            dis[e[i].v] = dis[now] + 1;
            q.push(e[i].v);
            xu.pb(e[i].v);
        }
    }
}

bool check() {
    xu.clear();
    int n, rt, sum = 0;
    ms(fa, 0);
    ms(b, -1);
    ms(dis,0);
    en = 0;
    cin >> n;
    rep(i, 1, n) {
        scanf("%d%d", &fa[i], &t[i]);
        if(fa[i] != 0)
            addedge(fa[i], i);
        sum += t[i];
        if(fa[i] == 0)
            rt = i;
    }
    if(sum % 3 != 0){return false;}
    sum /= 3;
    bfs(rt);

    repd(i, 1, n) {
        int now = xu[i-1];
        for(int j = b[now]; j!=-1; j = e[j].nxt)
            t[now] += t[e[j].v];
    }

    ms(dp, 0);
    repd(i, 1, n) {
        int now = xu[i-1];

        int u = 0;
        for(int j = b[now]; j != -1; j = e[j].nxt) {
            if(dp[e[j].v] != 0 && t[now] == sum * 2 && now != rt)
                return true;
            if(dp[e[j].v] != 0){
                if(u!=0){
                    return true;
                }
                else
                    u = dp[e[j].v];
            }
        }
        if(u != 0) dp[now] = u;
        if(t[now] == sum) dp[now] = now;
    }
    //rep(i,1,n)printf("%d ", dp[i]);
    return false;
}

int main() {
    freopen("strawberry.in", "r" ,stdin);
    freopen("strawberry.out", "w" ,stdout);
    int tc;
    cin >> tc;
    while(tc--) {
        if(check())printf("YES\n");
        else printf("NO\n");
    }
    return 0;
}