【BZOJ3727】Zadanie（PA2014 Final）-思维

测试地址：Zadanie
题目大意： 一棵树，第$i$个点有$a_i$个人，现在求出了$b_i$，为所有人走到点$i$的总路程，要求还原$a_i$。
做法： 本题需要用到思维。
在我们求$b_i$的时候，我们可以使用换根法，那么我们能不能用换根法，找到$a,b$之间的关系呢？
在换根法中，先随便选一个点作为根（这个根不是指换根法中的根，而是为了算法方便而求出的根），把根从点$i$换到父亲点$j$，$b$会增加$sum_i-(total-sum_i)=2sum_i-total$，其中$sum_i$为以$i$为根的子树内所有点的$a_i$之和，$total$为所有点的$a_i$之和。因此得到等式$b_j-b_i=2sum_i-total$。
这样类推下去，我们可以得到$n-1$个等式。但之中有$n$个未知数，所以我们还需要找到一个条件。我们注意到上面的等式中$b$都是差的形式，如果我们再多用一个点的$b$来连接和所有未知数之间的关系，应该就可以解了，这个点就是算法的根$root$。
我们发现，$b_{root}$就等于$\sum_{i\ne root}sum_i$。这可以通过一个简单的贡献变换得出。这样一来我们就有了$n$个等式，意味着这个方程组可以解了。
怎么解呢？首先当然是算出$total$。一开始的$n-1$个等式都可以化成下面的形式：$sum_i=\frac{total+b_{f(i)}-b_i}{2}$（其中$f(i)$表示$i$的父亲），于是有：
$2\sum_{i\ne root}sum_i=(n-1)total+\sum_{i\ne root}(b_{f(i)}-b_i)$
又因为$\sum_{i\ne root}sum_i=b_{root}$，所以$total$为：
$total=\frac{2b_{root}+\sum_{i\ne root}(b_i-b_{f(i)})}{n-1}$
这就可以很轻松地算出来了。进一步地，有了$total$后就可以通过前$n-1$个等式逐一算出每个点的$sum$了，实际上$total$就是根的$sum$。而从子树和$sum$恢复$a$也非常容易了，对于每个点$i$，对$sum_i$减去它所有儿子的$sum$即可得到$a_i$。这样我们就解决了这一题，时间复杂度为$O(n)$。
以下是本人代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n,first[300010]={0},tot=0,fa[300010]={0};
ll a[300010],b[300010],total;
struct edge
{
	int v,next;
}e[600010];

void insert(int a,int b)
{
	e[++tot].v=b;
	e[tot].next=first[a];
	first[a]=tot;
}

void solve(int v)
{
	total+=b[v]-b[fa[v]];
	for(int i=first[v];i;i=e[i].next)
		if (e[i].v!=fa[v])
		{
			fa[e[i].v]=v;
			solve(e[i].v);
		}
}

void finalsolve(int v)
{
	for(int i=first[v];i;i=e[i].next)
		if (e[i].v!=fa[v])
		{
			a[v]-=a[e[i].v];
			finalsolve(e[i].v);
		}
}

int main()
{
	scanf("%d",&n);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		insert(u,v),insert(v,u);
	}
	for(int i=1;i<=n;i++)
		scanf("%lld",&b[i]);
	
	total=b[1];
	solve(1);
	total/=(ll)(n-1);
	a[1]=total;
	for(int i=2;i<=n;i++)
		a[i]=(total-b[i]+b[fa[i]])>>1;
	finalsolve(1);
	for(int i=1;i<=n;i++)
		printf("%lld ",a[i]);
	
	return 0;
}