【题解】CF1681F-Unique Occurrences

给定一个 \(n\) 个结点的树，每条边有一个颜色，记 \(f(x,y)\) 表示结点 \(x\) 到 \(y\) 的路径上只出现了一次的颜色的数量，求 \(\sum\limits_{x < y}f(x, y)\) 。数据保证 \(n \le 5 \times 10^5\) 。

题面很短，解法很多。

由于我们是求所有 \(f(x,y)\) 的总和，根据经验我们将问题转化为对每个颜色 \(w\) 计算有多少路径经过恰好一条颜色为 \(w\) 的边。这样每种颜色是独立的，可以分开计算。

方法一：虚树

很自然我们可以想到虚树的解法，我们将颜色相同的边的两个端点保留下来，建立虚树后直接跑 DP。

方法二：二次扫描换根

如果我们设 DP 的状态 \(f_i\) 表示以 \(i\) 为根子树中，以 \(i\) 为起点，不经过颜色 \(w\) 的单链数量，同理 \(g_i\) 表示删去 \(i\) 的子树的数量，那么一条颜色为 \(w\) 的边的贡献就是两端 \(f,g\) 的乘积。\(f\) 可以在 DFS 的同时维护栈求得，然后二次扫描换根求得 \(g\)。

方法三：线段树分治

如果现在考虑颜色 \(w\)，那么删去所有该颜色的边，求得剩下每个连通块的大小，那么每条颜色 \(w\) 的边的贡献就是两端连通块大小的乘积。显然这个过程可以通过线段树分治快速维护。

方法四：线段树

固定一个根，如果我们把颜色为 \(w\) 的边的深度较深的点记为关键点，考虑用 DS 快速求出方法三中的连通块大小。

我们先将树上每个点赋值为 \(1\)，将关键点按 DFS 序排序，每次求得子树和记为连通块大小，然后将子树内所有点赋值为 \(0\)，这些都可以用数据结构快速维护。

---

代码实现的是方法三，因为好想好写，时间复杂度 \(\mathcal{O}(n\log ^2 n)\)。

#define N 500005
int n, fa[N], sz[N];
inline int get(int x){while(x != fa[x])x = fa[x]; return x;}
vector<Pr>u[N << 2], c[N];
#define ls (x << 1)
#define rs (ls | 1)
void ins(int x,int L,int R,int l,int r,Pr w){
	if(L >= l && R <= r){u[x].pb(w); return;}
	int mid = (L + R) >> 1;
	if(mid >= l)ins(ls, L, mid, l, r, w);
	if(mid < r)ins(rs, mid + 1, R, l, r, w);
}
LL ans;
void solve(int x,int l,int r){
	vector<int>dl;
	go(w, u[x]){
		int p = get(w.fi), q = get(w.se);
		if(sz[p] > sz[q])swap(p, q);
		dl.pb(p), sz[q] += sz[p], fa[p] = q;
	}
	if(l == r)go(w, c[l])ans += sz[get(w.fi)] * 1LL * sz[get(w.se)];
	else {
		int mid = (l + r) >> 1;
		solve(ls, l, mid), solve(rs, mid + 1, r);
	}
	reverse(dl.begin(), dl.end());
	go(w, dl)sz[fa[w]] -= sz[w], fa[w] = w;
}
int main() {
	read(n);
	rp(i, n - 1){
		int x, y, z; read(x, y, z);
		c[z].pb(mp(x, y));
		if(z > 1)ins(1, 1, n, 1, z - 1, {x, y});
		if(z < n)ins(1, 1, n, z + 1, n, {x, y});
	}
	rp(i, n)sz[fa[i] = i] = 1;
	solve(1, 1, n);
	printf("%lld\n", ans);
	return 0;
}