[POI2014] HOT-Hotels 加强版

[POI2014] HOT-Hotels

题面翻译

给定一棵树，在树上选 $3$ 个点，要求两两距离相等，求方案数。

题目描述

There are $n$ towns in Byteotia, connected with only $n-1$ roads.

Each road directly links two towns.

All the roads have the same length and are two way.

It is known that every town can be reached from every other town via a route consisting of one or more (direct-link) roads.

In other words, the road network forms a tree.

Byteasar, the king of Byteotia, wants three luxury hotels erected to attract tourists from all over the world.

The king desires that the hotels be in different towns and at the same distance one from each other.

Help the king out by writing a program that determines the number of possible locations of the hotel triplet in Byteotia.

输入格式

The first line of the standard input contains a single integer $n$ ($1\le n\le 10^5$), the number of towns in Byteotia.

The towns are numbered from $1$ to $n$.

The Byteotian road network is then described in $n-1$ lines.

Each line contains two integers $a$ and $b$ ($1\le a\le b\le n$) , separated by a single space, that indicate there is a direct road between the towns $a$ and $b$.

输出格式

The first and only line of the standard output should contain a single integer equal to the number of possible placements of the hotels.

样例 #1

样例输入 #1

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

样例输出 #1

5

最暴力的 $O(n^2)$ 应该是枚举一个点，搜索另外两个点

但是太暴力了，没有优化空间，所以要换一种方式。

定义 $dp_{i,j}$ 为子树 $i$ 内，距离 $i$ 为 $j$ 的有多少个

然后在背包合并的时候就可以知道 lca 为 $i$，距离为 $j$ 的点对有多少个。

还要再选一个点，定义 $g_{i,j}$ 为此时如果多一个距离 $i$ 为 $j$ 的点，答案会增加多少。合并子树时可以统计答案，顺便把 $f$ 合并到 $g$ 上。

#include<bits/stdc++.h>
using namespace std;
const int N=5005;
int n,fa[N],hd[N],e_num,rt;
long long f[N][N],g[N][N],ans;
struct edge{
	int v,nxt;
}e[N<<1];
void add_edge(int u,int v)
{
	e[++e_num]=(edge){v,hd[u]};
	hd[u]=e_num;
}
void dfs(int x,int y)
{
	f[x][0]++;
	for(int i=hd[x];i;i=e[i].nxt)
	{
		if(e[i].v==y)
			continue;
		dfs(e[i].v,x);
		for(int j=1;j<=n;j++)
		{
			ans+=g[x][j]*f[e[i].v][j-1];
			ans+=g[e[i].v][j+1]*f[x][j];
			g[x][j]+=f[x][j]*f[e[i].v][j-1];
			f[x][j]+=f[e[i].v][j-1];
		}
		for(int j=0;j<=n;j++)
			g[x][j]+=g[e[i].v][j+1];
	}
	ans+=g[x][0];
}
int main()
{
	scanf("%d",&n);
	for(int i=1,u,v;i<n;i++)
		scanf("%d%d",&u,&v),add_edge(u,v),add_edge(v,u);
	dfs(1,0);
	printf("%lld",ans);
}

然后发现 dp 数组的长度只和层数有关，所以可以对他进行长剖优化。

#include<bits/stdc++.h>
using namespace std;
const int N=5e5+5;;
int n,fa[N],hd[N],e_num,rt,df[N],dfn[N],dp[N],son[N],idx,to[N],in[N],tme;
long g[N<<1],ans,f[N];
struct edge{
	int v,nxt;
}e[N<<1];
void add_edge(int u,int v)
{
	e[++e_num]=(edge){v,hd[u]};
	hd[u]=e_num;
}
void sou(int x,int y)
{
	for(int i=hd[x];i;i=e[i].nxt)
	{
		if(e[i].v==y)
			continue;
		sou(e[i].v,x);
		if(dp[e[i].v]>dp[son[x]])
			son[x]=e[i].v;
	}
	dp[x]=dp[son[x]]+1;
}
void suo(int x,int y)
{
	dfn[x]=++idx;
	if(son[x])
		suo(son[x],x);
	for(int i=hd[x];i;i=e[i].nxt)
	{
		if(e[i].v==y||e[i].v==son[x])
			continue;
		suo(e[i].v,x);
	}
}
void doit(int x,int y)
{
	in[x]=++idx;
	if(son[x])
		doit(son[x],x);
	for(int i=hd[x];i;i=e[i].nxt)
		if(e[i].v^son[x]&&e[i].v^y)
			idx+=dp[e[i].v],doit(e[i].v,x);
}
void dfs(int x,int y)
{
	if(son[x])
		dfs(son[x],x);
	f[dfn[x]]=1;
	for(int i=hd[x];i;i=e[i].nxt)
	{
		if(e[i].v==y||e[i].v==son[x])
			continue;
		dfs(e[i].v,x);
		for(int j=1;j<=dp[e[i].v];j++)
		{
			ans+=g[in[x]-j]*f[dfn[e[i].v]+j-1];
			if(j+1<dp[e[i].v])
				ans+=g[in[e[i].v]-j-1]*f[dfn[x]+j];
			g[in[x]-j]+=1LL*f[dfn[x]+j]*f[dfn[e[i].v]+j-1];
			f[dfn[x]+j]+=f[dfn[e[i].v]+j-1];
		}
		for(int j=0;j<dp[e[i].v]-1;j++)
			g[in[x]-j]+=g[in[e[i].v]-j-1];
	}
	ans+=g[in[x]];
}
int main()
{
	scanf("%d",&n);
	for(int i=1,u,v;i<n;i++)
		scanf("%d%d",&u,&v),add_edge(u,v),add_edge(v,u);
	sou(1,0);
	suo(1,0);
	idx=dp[1];
	doit(1,0);
	dfs(1,0);
	printf("%lld",ans);
}