题目大意:给定一棵 N 个节点的树,点有点权,问是否可以将整棵树分成三个权值和均相等的联通块。
题解:
直接模拟整个过程即可,记录 f[u] 表示以 u 为根的子树中的权值和,若 f[u] 为 sum/3,则割断这棵树,在代码中的具体表现为将 f[u] 置零,即:将对 u 的父节点的贡献消去。
代码如下
#include <bits/stdc++.h>
#define pb push_back
using namespace std;
const int maxn=1e6+10;
vector<int> G[maxn],ans;
int n,rt,sum,a[maxn],f[maxn];
void dfs(int u){
f[u]=a[u];
for(auto v:G[u]){
dfs(v);
f[u]+=f[v];
}
if(f[u]==sum)ans.pb(u),f[u]=0;
}
void read_and_parse(){
scanf("%d",&n);
for(int i=1,fa;i<=n;i++){
scanf("%d%d",&fa,&a[i]);
sum+=a[i];
if(!fa)rt=i;
else G[fa].pb(i);
}
}
void solve(){
if(sum%3!=0)return (void)puts("-1");
sum/=3;
dfs(rt);
if(ans.size()<3)puts("-1");
else printf("%d %d
",ans[0],ans[1]);
}
int main(){
read_and_parse();
solve();
return 0;
}