Educational Codeforces Round 129 F. Unique Occurrences

传送门
\(\texttt{difficulty:2300}\)

题目大意

一棵 \(n(2\le n\le 5\cdot 10^5)\) 个节点的树，每条边有权值 \(x(1\le x\le n)\) ， \(f(u,v)\) 表示从 \(u\) 到 \(v\) 路径上所有仅出现过 \(1\) 次的数字的个数，求 \(\sum_{1\le u<v\le n}f(u,v)\) 。

思路

我们考虑每一条边对于答案的贡献，其贡献为所有路径仅经过该权值的边一次的节点对数，我们如果删掉所有权值为 \(x\) 的边，那么原来的树就会变为若干连通块，连通块之间可以通过被删除的边重新连成一棵树，我们记为 \(t_x\) ，那么每条权值为 \(x\) 的边对答案的贡献就是就是其连接的连通块的大小的乘积。
我们可以通过 \(\texttt{dfs}\) 来解决，先求出以每个节点为根的子树的大小 \(siz_v\) 。之后再进行一遍 \(\texttt{dfs}\) ，我们对每种权值 \(x\) 分别建立 \(t_x\) ，同时建立 \(n\) 个新的节点作为这些树的根（不然的话所有的 \(t_x\) 都会共用一个根），并且让它们的 \(siz\) 都为 \(n\) 。我们可以用所有向上的边的权值为 \(x\) 的节点表示删除权值为 \(x\) 的边后的连通块，将该节点加入到 \(t_x\) 中，同时我们再对每种权值用栈来记录其当前的父亲节点来进行连边。
对于连通块大小的计算，我们在新加入一个节点时，其连通块的大小就设为 \(siz_v\) ，同时将其父亲节点所代表的连通块大小减去 \(siz_v\) ，这样就可以得到每个连通块的大小，最后对所有 \(t_x\) 进行 \(\texttt{dfs}\) 求出答案即可，原树中的每个节点仅会加入到一个 \(t_x\) 当中，于是复杂度为 \(O(n)\) 。

代码

#include<bits/stdc++.h>
#include<unordered_map>
#include<unordered_set>
using namespace std;
using LL = long long;
using ULL = unsigned long long;
using PII = pair<LL, LL>;
using TP = tuple<int, int, int>;
#define all(x) x.begin(),x.end()
#define mk make_pair
//#define int LL
//#define lc p*2
//#define rc p*2+1
#define endl '\n'
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f
#pragma warning(disable : 4996)
#define IOS ios::sync_with_stdio(0),cin.tie(0),cout.tie(0)
const double eps = 1e-8;
const LL MOD = 1000000007;
const LL mod = 998244353;
const int maxn = 1000010;

struct edge {
	int to, cost;
};
vector<edge>G[maxn];
vector<int>T[maxn];
stack<int>S[maxn];
LL N, siz[maxn], nsiz[maxn], root[maxn], ans = 0;

void add_edge(int from, int to, int cost)
{
	G[from].push_back(edge{ to,cost });
	G[to].push_back(edge{ from,cost });
}

void add_edge(int from, int to)
{
	T[from].push_back(to);
	T[to].push_back(from);
}

void dfs1(int v, int p)
{
	siz[v] = 1;
	for (auto& [to, c] : G[v])
	{
		if (to == p)
			continue;
		dfs1(to, v);
		siz[v] += siz[to];
	}
}

void dfs2(int v, int p, int x)
{
	if (x)
	{
		nsiz[v] = siz[v];
		add_edge(S[x].top(), v);
		nsiz[S[x].top()] -= siz[v];
		S[x].push(v);
	}
	for (auto& [to, c] : G[v])
	{
		if (to == p)
			continue;
		dfs2(to, v, c);
	}
	if (x)
		S[x].pop();
}

void dfs3(int v, int p)
{
	for (auto&to : T[v])
	{
		if (to == p)
			continue;
		ans += nsiz[v] * nsiz[to];
		dfs3(to, v);
	}
}

void solve()
{
	for (int i = 1; i <= N; i++)
		root[i] = N + i, S[i].push(N + i), nsiz[root[i]] = N;
	dfs1(1, 0), dfs2(1, 0, 0);
	for (int i = 1; i <= N; i++)
		dfs3(root[i], 0);
	cout << ans << endl;
}

int main()
{
	IOS;
	cin >> N;
	int u, v, x;
	for (int i = 1; i < N; i++)
		cin >> u >> v >> x, add_edge(u, v, x);
	solve();

	return 0;
}