【BZOJ4543】[POI2014]Hotel加強版 長鏈剖分+DP

【BZOJ4543】[POI2014]Hotel加強版

Description

同OJ3522
數據范圍：n<=100000

Sample Input

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

Sample Output

5

題解：很神的做法。

用f[x][a]表示x子樹中有多少個深度為a的點，g[x][a]表示x子樹中有多少到lca距離=d的點對，且lca的深度為d-a。那么容易得到轉移方程：

f[x][a]+=f[y][a-1]
g[x][a]+=g[y][a+1]+f[x][a]*f[y][a]
ans+=f[x][a]*g[y][a+1]+g[x][a]*f[y][a-1]

（以上方程的真正順序並不是給出的順序）

顯然一次轉移的復雜度是maxdep級別的，總復雜度是$O(n^2)$，如何優化呢？

發現每個節點的第一次轉移就相當於將y的f數組向右移一格，g數組向左移一格。所以考慮用指針優化這個過程，這樣每個點的第一次轉移就是O(1)的了。顯然我們應該選擇maxdep最大的兒子去優化，即：我們對原樹進行長鏈剖分，對於同一條鏈上的，用指針O(1)轉移，其余的暴力轉移。

復雜度是。。。O(n)?表示不太會證。

#include <cstdio>
#include <cstring>
#include <iostream>
using namespace std;
const int maxn=100010;
typedef long long ll;
int n,cnt;
ll ans;
int to[maxn<<1],next[maxn<<1],head[maxn],md[maxn],son[maxn],fa[maxn];
ll mem[1000010];
ll *f[maxn],*g[maxn],*now=mem+1;
inline void add(int a,int b)
{
	to[cnt]=b,next[cnt]=head[a],head[a]=cnt++;
}
void dfs1(int x)
{
	md[x]=0;
	for(int i=head[x];i!=-1;i=next[i])	if(to[i]!=fa[x])
	{
		fa[to[i]]=x,dfs1(to[i]),md[x]=max(md[x],md[to[i]]+1);
		if(md[to[i]]>md[son[x]])	son[x]=to[i];
	}
}
void dfs2(int x)
{
	int i,j,y;
	if(son[x])
	{
		f[son[x]]=f[x]+1,g[son[x]]=g[x]-1;
		dfs2(son[x]);
	}
	f[x][0]=1,ans+=g[x][0];
	for(i=head[x];i!=-1;i=next[i])	if(to[i]!=fa[x]&&to[i]!=son[x])
	{
		y=to[i],f[y]=now,now+=md[y]+1,g[y]=now+md[y]+1,now+=md[y]*2+2,dfs2(y);
		for(j=md[y];j>=0;j--)
		{
			if(j)	ans+=f[x][j-1]*g[y][j];
			ans+=g[x][j+1]*f[y][j];
			g[x][j+1]+=f[x][j+1]*f[y][j];
		}
		for(j=0;j<=md[y];j++)
		{
			if(j)	g[x][j-1]+=g[y][j];
			f[x][j+1]+=f[y][j];
		}
	}
}
inline int rd()
{
	int ret=0,f=1;	char gc=getchar();
	while(gc<'0'||gc>'9')	{if(gc=='-')	f=-f;	gc=getchar();}
	while(gc>='0'&&gc<='9')	ret=ret*10+gc-'0',gc=getchar();
	return ret*f;
}
int main()
{
	n=rd(),md[0]=-1;
	int i,a,b;
	memset(head,-1,sizeof(head));
	for(i=1;i<n;i++)	a=rd(),b=rd(),add(a,b),add(b,a);
	dfs1(1);
	f[1]=now,now+=md[1]+1,g[1]=now+md[1]+1,now+=md[1]*2+2,dfs2(1);
	printf("%lld",ans);
	return 0;
}