CF767C 经典的树形DP

Posted 2021-02-06 kenlig

n个节点的树
第i个节点权值为a[i]
n<=10^6
-100<=a[i]<=100
问是否能够删除掉两条边,使得该树分成三个不为空,并且每部分权值之和相等.
无解输出-1 否则输出要删除边(u->v)的v节点序号.

以上是题面。如果觉得这个题面看不懂的话可以去CodeForces现场看一下（逃）
下面是解法。
考虑存在这么三个子树，那么这个树肯定就能分成这么三个子树，而这个树的点权之和肯定能平均分成三个子树，那么这个树的点权 mod 3 肯定为0.
这是一个显然的剪枝。
言归正传，我们来说说正解。首先这个题的解法当然是树形DP一贯的解法：推。我们从下往上更新，就可以更新到这个子树了。如果我们从下往上获得了一个子树，并且这个子树的权值是树的权值的1/3，那么显然这个子树是符合条件的，切一刀，我们就得到了子树。代码中cnt是切的次数。
ksum这个数组就是子树的大小。
那么本题的解法已经很显然了，就那么遍历就完事了。
代码有注释，方便各位理解。
C++：

#include<iostream>
#include<cstdio>
#include<cmath>
#include<algorithm>
#define LL long long
#define maxn 1000001
using namespace std;
struct Edge{
    int nxt,to,from;
}edge[maxn*2];//因为这是个树 
int num_edge,head[maxn],a[maxn],n,fa,total,sum,cnt,ans[maxn],ksum[maxn];
inline int qread(){
    char ch=getchar();
    int f=1,x=0;
    while(ch>‘9‘||ch<‘0‘){if(ch==‘-‘) f=-1;ch=getchar();}
    while(ch>=‘0‘&&ch<=‘9‘){x=x*10+ch-‘0‘;ch=getchar();}
    return x*f;
}
inline void addedge(int from,int to){ //加上from的原因是我们要输出他的from 
    edge[++num_edge].nxt=head[from];
    edge[num_edge].from=from;
    edge[num_edge].to=to;
    head[from]=num_edge;
} 
int dfs(int x,int fath){
    ksum[x]=a[x];
    for(int i=head[x];i;i=edge[i].nxt){//找祖先 
        int y=edge[i].to;
        if(y!=fath){
            dfs(y,x);
            ksum[x]+=ksum[y];
        }
    }
    if(ksum[x]==sum){ //当满足了之后 
        cnt++;
        ans[cnt]=x;
        ksum[x]=0;
    }
}
int main(){
    n=qread();//快读，否则的话大数据卡死 
    for(int i=1;i<=n;i++){
        int x;
        x=qread();
        a[i]=qread();
        if(x!=0) {
            addedge(i,x);
            addedge(x,i);
        }
        else fa=i;
        total+=a[i];
    }
    if(total%3!=0) printf("-1
");
    else{
        sum=total/3;
        dfs(fa,0);
        if(cnt>=3) printf("%d %d
",ans[2],ans[1]);
        else printf("-1
");
    }
    return 0;
}