[bzoj3727]PA2014 Final Zadanie

题目大意

一颗n个节点的树。
$bi=\sum_{j=1}^naj*dis(j,i)$
给定b，请求出a。

好题

设size表示子树的a值和。设x表示整颗树a值和。
那么对于i是j的父亲，有
$b[i]-size[j]+x-size[j]=b[j]$
$x-2*size[j]=b[j]-b[i]=c[j]$
我们可以把每个a都表示出来（a[1]不能表示，我们以1为根）
$a[i]=size[i]-\sum_{j是i儿子} size[j]$
$2a[i]=-c[i]+\sum_{j是i儿子}c[j]-(d[i]-1)*x$
d[i]表示i的儿子数，然后为了方便设$e[i]=-c[i]+\sum_{j是i儿子}c[j]$
接着观察
$b[1]=\sum_{i=2}^na[i]*dep[i]$
dep[i]表示深度
那么代入a后，得到关于x的方程，可以解出x，也就是size[1]。
然后可以把所有size以及a均解出。

#include<cstdio>
#include<algorithm>
#define fo(i,a,b) for(i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const int maxn=300000+10;
int a[maxn],b[maxn],c[maxn],d[maxn],dep[maxn],size[maxn];
ll e[maxn];
int h[maxn],go[maxn*2],next[maxn*2];
int i,j,k,l,t,n,m,tot;
ll sum,num,cnt;
void add(int x,int y){
    go[++tot]=y;
    next[tot]=h[x];
    h[x]=tot;
}
void dfs(int x,int y){
    if (x!=1){
        c[x]=b[x]-b[y];
        e[x]=-c[x];
        dep[x]=dep[y]+1;
    }
    int t=h[x];
    while (t){
        if (go[t]!=y){
            d[x]++;
            dfs(go[t],x);
            e[x]+=(ll)c[go[t]];
        }
        t=next[t];
    }
}
void dg(int x,int y){
    if (x!=1) size[x]=(sum-c[x])/2;
    int t=h[x];
    a[x]=size[x];
    while (t){
        if (go[t]!=y){
            dg(go[t],x);
            a[x]-=size[go[t]];
        }
        t=next[t];
    }
}
int main(){
    scanf("%d",&n);
    fo(i,1,n-1){
        scanf("%d%d",&j,&k);
        add(j,k);add(k,j);
    }
    fo(i,1,n) scanf("%d",&b[i]);
    dfs(1,0);
    fo(i,2,n) num+=(ll)e[i]*dep[i];
    num=(ll)2*b[1]-num;
    fo(i,2,n) cnt+=(ll)(d[i]-1)*dep[i];
    cnt=-cnt;
    sum=num/cnt;
    size[1]=sum;
    dg(1,0);
    fo(i,1,n-1) printf("%d ",a[i]);
    printf("%d\n",a[n]);
}