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

\(\mathcal{Description}\)

  Link.

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

  \(n\le10^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\n", ans );
	return 0;
}