Unique Occurrences
Contest: Educational Codeforces Round 129 (Rated for Div. 2)
Difficulty rating: 2300
Tags: data structures, dfs and similar, divide and conquer, dp, dsu, trees

You are given a tree, consisting of $$$n$$$ vertices. Each edge has an integer value written on it.

Let $$$f(v, u)$$$ be the number of values that appear exactly once on the edges of a simple path between vertices $$$v$$$ and $$$u$$$.

Calculate the sum of $$$f(v, u)$$$ over all pairs of vertices $$$v$$$ and $$$u$$$ such that $$$1 \le v < u \le n$$$.

Input format:
The first line contains a single integer $$$n$$$ ($$$2 \le n \le 5 \cdot 10^5$$$) — the number of vertices in the tree.

Each of the next $$$n-1$$$ lines contains three integers $$$v, u$$$ and $$$x$$$ ($$$1 \le v, u, x \le n$$$) — the description of an edge: the vertices it connects and the value written on it.

The given edges form a tree.

Output format:
Print a single integer — the sum of $$$f(v, u)$$$ over all pairs of vertices $$$v$$$ and $$$u$$$ such that $$$v < u$$$.

Constraints:
Time limit: 6.0 s
Memory limit: 1024.0 MB

Example input:
3
1 2 1
1 3 2

Example output:
4

Example input:
3
1 2 2
1 3 2

Example output:
2

Editorial:
Let's use contribution to the sum technique to simplify the problem. Instead of counting the number of colors that occure only once for each path, let's, for each color, count the number of paths that contain this color exactly once. Now we can solve the problem independently for each color, and sum up the answers.
The first intended solution was the following. So we want to calculate the answer for some color $$$c$$$. Mark all edges of color $$$c$$$ as good, the rest are bad. Then we can calculate $$$dp_{v,i}$$$ — the number of paths up to vertex $$$v$$$ such that they contain either $$$0$$$ or $$$1$$$ good edges. The transitions should be pretty easy, and the answer should be updated when you consider gluing up paths from different children in each vertex. Obviously, this is $$$O(n)$$$ per color, so $$$O(n^2)$$$ overall.
However, we can only calculate this dynamic programming as easily on a virtual tree of vertices adjacent to all good edges. How to calculate the dp for some vertex $$$v$$$? First, push the paths from all virtual children to $$$v$$$. That was enough in the dp for the entire tree but now there are also removed vertices that could also have paths starting in them. All these paths contain $$$0$$$ good edges (otherwise, they would have had virtual vertices on them). Their amount is the following: the size of the real subtree of $$$v$$$ minus the sizes of real subtrees of all its virtual children. The rest is exactly the same as in the dp on the real tree.
A little fun trick. Usually, you want to add lca of adjacent vertices to the virtual tree. But that's actually not needed here: you can just add the root of the tree and link the vertices without a parent to them. That won't change the result of the dp.
That solution works in $$$O(n \log n)$$$ or $$$O(n)$$$.
The second intended solution is slower complexity-wise but not time-wise. In the first solution we wanted to leave only the good edges in the tree. Here, we want to remove only them. Consider the resulting connected components. What's the number of paths that contain only one of the good edges? It's actually the product of sizes of the connected components this edge connects.
So we want to remove edges, add edges and maintain the sizes of the connected components of the tree. That's basically the same problem as dynamic connectivity. The $$$O(n \log^2 n)$$$ implementation works well enough.