【换根DP例题】——P3047 [USACO12FEB] Nearby Cows G

---

题目

原题链接

描述

Farmer John has noticed that his cows often move between nearby fields. Taking this into account, he wants to plant enough grass in each of his fields not only for the cows situated initially in that field, but also for cows visiting from nearby fields.
Specifically, FJ’s farm consists of N fields (1 <= N <= 100,000), where some pairs of fields are connected with bi-directional trails (N-1 of them in total). FJ has designed the farm so that between any two fields i and j, there is a unique path made up of trails connecting between i and j. Field i is home to C(i) cows, although cows sometimes move to a different field by crossing up to K trails (1 <= K <= 20).
FJ wants to plant enough grass in each field i to feed the maximum number of cows, M(i), that could possibly end up in that field – that is, the number of cows that can potentially reach field i by following at most K trails. Given the structure of FJ’s farm and the value of C(i) for each field i, please help FJ compute M(i) for every field i.

给你一棵 n 个点的树，点带权，对于每个节点求出距离它不超过 k 的所有节点权值和 m i。

输入格式

• Line 1: Two space-separated integers, N and K.

• Lines 2…N: Each line contains two space-separated integers, i and j ($1<=i,j<=N$) indicating that fields i and j are directly connected by a trail.

• Lines N+1…2N: Line N+i contains the integer C(i). ($0<=C(i)<=1000$)

• 第一行两个正整数n,k。

• 接下来n−1 行，每行两个正整数 u,v，表示 u,v 之间有一条边。

• 最后 n 行，每行一个非负整数 c i，表示点权。

输出格式

• Lines 1…N: Line i should contain the value of M(i).

• 输出 n 行，第 i 行一个整数表示 m i。

输入/输出例子

输入

6 2
5 1
3 6
2 4
2 1
3 2
1
2
3
4
5
6

输出

15
21
16
10
8
11

提示

【数据范围】
对于 100%的数据：$1\le n\le 100000$ ，$1\le k\le 20$，$0\le c_i\le 1000$

---

解题思路

分析

考虑暴力方法：每个点都进行一次树形DP，时间复杂度$O(n^2)$
优化：对于根节点切换时，我们需要用$O(1)$的时间复杂度算出另一个根的数据，所以用换根DP

初始化/第一遍$dfs$

定义：
$f[i][j]$表示以$i$为根节点时整棵树中距离$i$不超过$j$的节点数
$g[i][j]$表示以$i$为根节点的子树中距离$i$为$j$的节点数。
所以我们可以得到：
$g[u][k]+=g[to][k-1]$

code1

void dfs(int u,int fa)
{
	for(int i=head[u];i;i=edge[i].next)
	{
		int w=edge[i].w,to=edge[i].v;
		if(to==fa)
			continue;
		dfs(to,u);
		foru(j,1,k)
			g[u][j]+=g[to][j-1];
	}
}

$dfs$结束后，我们再对$g$做前缀和，那么现在$g[i][j]$就表示以$i$为根节点的子树中距离$i$不超过$j$的节点个数
同时，由于我们第一遍$dfs$是以$1$为根节点的，所以: $f[1][i]=g[1][i]$

code2

	f[1][0]=g[1][0];
	foru(i,1,n)
	{
		foru(j,1,20)
			g[i][j]+=g[i][j-1];
		f[1][i]=g[1][i];
	}

动态转移方程/第二遍$dfs$

我们已经知道了$f[1][i]$，接下里需要转移$f[j][i]$

如图，$f[v][k]=g[v][k]+A$
（A代表A中距离v不超过k的点的点数）
而$A=f[j][k-1]-g[v][k-2]$
注：这个可以自己推一下

code

void dfs_(int u,int fa)
{
	for(int i=head[u];i;i=edge[i].next)
	{
		int w=edge[i].w,to=edge[i].v;
		if(to==fa)
			continue;
		f[to][0]=g[to][0];
		f[to][1]=g[to][1]+f[u][0];
		foru(j,2,k)
			f[to][j]=g[to][j]+f[u][j-1]-g[to][j-2];
		dfs_(to,u);
	}
}

AC Code

#include<bits/stdc++.h>
#define mod 1
#define int long long
#define foru(i,a,b) for(register int i=a;i<=b;i++)
#define ford(i,a,b) for(register int i=a;i>=b;i--)
using namespace std;
const int N=1e5+1;
int n,k,g[N][21],f[N][21],head[N],idx,x,y;
struct fy
{
	int v,w,next;
}edge[N<<1];
void add(int x,int y,int z)
{
	edge[++idx].v=y,edge[idx].w=z,edge[idx].next=head[x],head[x]=idx;
}
inline char gc()
{
  	static char now[1<<20],*S,*T;
  	if(T==S)
  	{
		T=(S=now)+fread(now,1,1<<20,stdin);
    	if(T==S) 
			return EOF;
  	}
  return *S++;
}
template <typename T>
inline void Read(T&x) 
{
  	x=0;
  	char c=gc();
  	while(c<'0'||c>'9') 
  		c=gc();
  	x=c-'0';
  	while((c=gc())>='0'&&c<='9') 
  		x=x*10+c-'0';
}
template <typename T, typename... Args>
inline void Read(T&x,Args&...args)
{
  	Read(x);
  	Read(args...);
}
void dfs(int u,int fa)
{
	for(int i=head[u];i;i=edge[i].next)
	{
		int w=edge[i].w,to=edge[i].v;
		if(to==fa)
			continue;
		dfs(to,u);
		foru(j,1,k)
			g[u][j]+=g[to][j-1];
	}
}
void dfs_(int u,int fa)
{
	for(int i=head[u];i;i=edge[i].next)
	{
		int w=edge[i].w,to=edge[i].v;
		if(to==fa)
			continue;
		f[to][0]=g[to][0];
		f[to][1]=g[to][1]+f[u][0];
		foru(j,2,k)
			f[to][j]=g[to][j]+f[u][j-1]-g[to][j-2];
		dfs_(to,u);
	}
}
signed main()
{
	Read(n,k);
	foru(i,2,n)
		Read(x,y),add(x,y,1),add(y,x,1);
	foru(i,1,n)
		Read(g[i][0]);
	dfs(1,0);
	f[1][0]=g[1][0];
	foru(i,1,n)
	{
		foru(j,1,20)
			g[i][j]+=g[i][j-1];
		f[1][i]=g[1][i];
	}
	dfs_(1,0);
	foru(i,1,n)
		printf("%lld\n",f[i][k]);
}

更多方法

更多方法