Solution -「POI 2014」「洛谷 P5904」HOT-Hotels 加强版

(mathcal{Description})

  Link.

  给定一棵 (n) 个点的树，求无序三元组 ((u,v,w)) 的个数，满足其中任意两点树上距离相等。

  (nle10^5)。

(mathcal{Solution})

  考虑如何计数。对于任意三元组 ((u,v,w))，我们仅在其两两路径所进过的树上最高点对其统计一次。如图：

  对于三元组 ((4,6,7))，我们仅希望在 (1) 处统计它的贡献。

  考虑 DP，记 (d(u,v)) 表示 (u) 到 (v) 的树上距离。令 (f(u,i)) 表示 (u) 子树内 (v) 的个数，满足 (d(u,v)=i)；(g(u,i)) 表示 (u) 子树内无序二元组 ((p,q)) 的个数，满足 (d(p,operatorname{lca}(p,q))=d(q,operatorname{lca}(p,q))=d(operatorname{lca}(p,q),u)+i)。例如上图的 (g(2,2)=1)，无序二元组 ((4,6)) 满足条件。

  如此设计状态的意义在于，在 (g(u,i)) 的基础上，在 (u) 子树的外部接上一个满足 (d(u,w)=i) 的 (w) 就能构成三元组，并且三元组恰好在最高点 (u) 计数。

  暴力转移比较显然，发现状态的第二维的范围均为 (u) 向下的最长链长，所以用长链剖分优化，直接移动指针继承 (mathcal O(1)) 继承长儿子信息，做到均摊 (mathcal O(n)) 转移，故总复杂度 (mathcal O(n))。

(mathcal{Code})

#include <cstdio>

#define alloc( u ) 
	( f[u] = cur, g[u] = ( cur += dep[u] << 1 ), cur += dep[u] << 1 )

typedef long long LL;

const int MAXN = 1e5;
int n, ecnt, head[MAXN + 5];
int dep[MAXN + 5], son[MAXN + 5];
LL ans, *f[MAXN + 5], *g[MAXN + 5], pool[MAXN * 5], *cur = pool;

struct Edge { int to, nxt; } graph[MAXN * 2 + 5];

inline void link ( const int s, const int t ) {
	graph[++ ecnt] = { t, head[s] };
	head[s] = ecnt;
}

inline void init ( const int u, const int fa ) {
	for ( int i = head[u], v; i; i = graph[i].nxt ) {
		if ( ( v = graph[i].to ) ^ fa ) {
			init ( v, u );
			if ( dep[v] > dep[son[u]] ) son[u] = v;
		}
	}
	dep[u] = dep[son[u]] + 1;
}

inline void solve ( const int u, const int fa ) {
	if ( son[u] ) {
		f[son[u]] = f[u] + 1, g[son[u]] = g[u] - 1;
		solve ( son[u], u );
	}
	f[u][0] = 1, ans += g[u][0];
	for ( int i = head[u], v; i; i = graph[i].nxt ) {
		if ( ( v = graph[i].to ) ^ fa && v ^ son[u] ) {
			alloc ( v ), solve ( v, u );
			for ( int j = 0; j < dep[v]; ++ j ) {
				if ( j ) ans += f[u][j - 1] * g[v][j];
				ans += g[u][j + 1] * f[v][j];
			}
			for ( int j = 0; j < dep[v]; ++ j ) {
				g[u][j + 1] += f[u][j + 1] * f[v][j];
				if ( j ) g[u][j - 1] += g[v][j];
				f[u][j + 1] += f[v][j];
			}
		}
	}
}

int main () {
	scanf ( "%d", &n );
	for ( int i = 1, u, v; i < n; ++ i ) {
		scanf ( "%d %d", &u, &v );
		link ( u, v ), link ( v, u );
	}
	init ( 1, 0 );
	alloc ( 1 );
	solve ( 1, 0 );
	printf ( "%lld
", ans );
	return 0;
}