bzoj 3727: PA2014 Final Zadanie 树形dp

题意

吉丽YY了一道神题，题面是这样的：
“一棵n个点的树，每条边长度为1，第i个结点居住着a[i]个人。假设在i结点举行会议，所有人都从原住址沿着最短路径来到i结点，行走的总路程为b[i]。输出所有b[i]。”
吉丽已经造好了数据，但熊孩子把输入文件中所有a[i]给删掉了。你能帮他恢复吗？
2<=n<=300000

分析

挺有趣的一道题。
先把1设为根，设size[x]表示x的子树中有多少个人，s表示整棵树的人数，显然有s=size[1]。
对于每条边(fa,x)，不难列出一个方程，b[x]+size[x]-s+size[x]=b[fa]
画一下可以得到s=2*size[x]+b[x]-b[fa]
这样我们就可以得到n个未知数n-1个方程的一个方程组。
要想得出解显然还需要一个方程，还那还有一个方程是什么呢？
考虑如何求b[1]，不难发现b[1]=size[2]+size[3]+…+size[n]。
这样我们就凑齐了n个方程。
对于由边得到的n-1个方程，将其变成size[x]=(s-b[x]+b[fa])/2，然后代进最后一个方程中，发现除了s以外的未知数都被消掉了，那么就可以得出s，之后每个点的size也都可以求出来了。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
using namespace std;

typedef long long LL;

const int N=300005;

int n,cnt,last[N];
LL b[N],s,size[N];
struct edge{int to,next;}e[N*2];

void addedge(int u,int v)
{
    e[++cnt].to=v;e[cnt].next=last[u];last[u]=cnt;
    e[++cnt].to=u;e[cnt].next=last[v];last[v]=cnt;
}

void dfs1(int x,int fa)
{
    if (fa) s-=b[fa]-b[x];
    for (int i=last[x];i;i=e[i].next)
        if (e[i].to!=fa) dfs1(e[i].to,x);
}

void dfs2(int x,int fa)
{
    if (fa) size[x]=(b[fa]-b[x]+s)/2;
    if (fa) size[fa]-=size[x];
    for (int i=last[x];i;i=e[i].next)
        if (e[i].to!=fa) dfs2(e[i].to,x);
}

int main()
{
    scanf("%d",&n);
    for (int i=1;i<n;i++)
    {
        int x,y;scanf("%d%d",&x,&y);
        addedge(x,y);
    }
    for (int i=1;i<=n;i++) scanf("%lld",&b[i]);
    s=b[1]*2;
    dfs1(1,0);
    s/=n-1;size[1]=s;
    dfs2(1,0);
    for (int i=1;i<=n;i++) printf("%lld ",size[i]);
    return 0;
}