设sum是所有灯泡的亮度之和
有两种情况:
一种是存在结点U和V,U是V的祖先,并且U的子树权值和为sum/3*2,且U不是根,且V的子树权值和为sum/3。
另一种是存在结点U和V,他们之间没有祖先关系,两者的子树权值和都是sum/3。(已经出栈的结点和当前访问的结点之间,必然没有祖先关系)
两次dfs解决。
#include<cstdio> #include<algorithm> #include<cstdlib> using namespace std; int n,a[1000010],fa[1000010],b[1000010]; int v[1000010],next[1000010],first[1000010],e,root,sum; void AddEdge(int U,int V) { v[++e]=V; next[e]=first[U]; first[U]=e; } void dfs(int U) { b[U]=a[U]; for(int i=first[U];i;i=next[i]) { dfs(v[i]); b[U]+=b[v[i]]; } } void df2(int U,int ance) { if(ance && b[U]==sum/3) { printf("%d %d ",min(ance,U),max(ance,U)); exit(0); } for(int i=first[U];i;i=next[i]) if(ance) df2(v[i],ance); else if(U!=root && b[U]==sum/3*2) df2(v[i],U); else df2(v[i],0); } int other; void df3(int U) { if(b[U]==sum/3 && other) { printf("%d %d ",min(other,U),max(other,U)); exit(0); } for(int i=first[U];i;i=next[i]) df3(v[i]); if(U!=root && b[U]==sum/3) other=U; } int main() { // freopen("c.in","r",stdin); int x; scanf("%d",&n); for(int i=1;i<=n;++i) { scanf("%d%d",&x,&a[i]); sum+=a[i]; if(x) { fa[i]=x; AddEdge(x,i); } else root=i; } if(sum%3!=0) { puts("-1"); return 0; } dfs(root); df2(root,0); df3(root); puts("-1"); return 0; }