Domain: Algorithms & Data Structures

You are given a tree with `n` vertices (`1` to `n`). Vertex `v` has color `c_v ∈ [1..n]`.

A simple path is **beautiful** if:
1. it has at least 2 vertices;
2. its endpoints have the same color;
3. no other vertex on the path has that endpoint color.

Paths are **undirected**.

### Input
- `n` (`2 ≤ n ≤ 2·10^5`)
- colors `c_1, ..., c_n`
- `n-1` edges

### Output
Output the number of beautiful simple paths.

Grading rubric:
criterion: Does the answer define a rooted tree (choose any root, e.g. 1) and process vertices in postorder (children before parent)?
weight: 10
criterion: Does the answer define cnt[v][x] as the number of vertices of color x in subtree(v) such that the path from v to that vertex contains no other vertex of color x (i.e., x-top-level w.r.t. v)?
weight: 10
criterion: Does the answer state that for any vertex v and any color x ≠ c[v], the number of beautiful paths whose LCA is v and whose endpoints have color x equals the sum over pairs of distinct children subtrees: (top-level x count in child A) × (top-level x count in child B)?
weight: 20
criterion: Does the answer state that for x = c[v], no beautiful path can have endpoints in two different child subtrees with LCA=v (because v would be an internal vertex of color x), so only paths (v, w) are allowed for x=c[v]?
weight: 10
criterion: Does the answer include the exact counting rule at vertex v for color c[v] as: add totalCv to the answer, where totalCv is the total number of c[v]-top-level vertices across all children of v?
weight: 10
criterion: Does the answer update the carried map at v by setting cnt[v][c[v]] = 1 (representing vertex v itself) and not propagating any descendant top-level vertices of color c[v] above v?
weight: 10
criterion: Does the answer describe a small-to-large (DSU-on-tree) merge where the largest child map is reused and all other child maps are merged into it, adding cross-pair contributions as ans += big[x] * small[x] for each x ≠ c[v] during the merge?
weight: 10
criterion: Does the answer state that each undirected beautiful path is counted exactly once by counting it at the LCA of its endpoints (or an equivalent unique-counting argument)?
weight: 10
criterion: Does the answer state a valid complexity bound for the DSU-on-tree implementation, e.g. O(n log^2 n) with ordered maps (or expected O(n log n) with hash maps), for n ≤ 2·10^5?
weight: 10

Answer:
### Expert Solution

Root the tree arbitrarily (e.g. at vertex 1) and process vertices in postorder.

#### 1) Top-level counts
For each vertex `v`, maintain a map `cnt[v]` where:

`cnt[v][x] = #{ w ∈ subtree(v) : c_w = x and there is no other x on the path v → w }`

So `w` is **`x`-top-level w.r.t. `v`**.

This “top-level” property ensures that once a vertex of color `x` is encountered, deeper vertices of color `x` do not propagate above it.

#### 2) Unique counting by LCA
Count each undirected path exactly once at the **LCA** of its endpoints (in the chosen rooting).

Let `v` have children `u_1, ..., u_k`.

#### 3) Counting contributions at vertex v

- For any color `x ≠ c[v]`:
  choosing endpoints of color `x` from **two different child subtrees** creates a beautiful path with LCA `v`, because `v` is on the path but has color `≠ x`.
  The number of such paths is:
  `∑_{A<B} cnt[u_A][x] · cnt[u_B][x]`

- For color `x = c[v]`:
  endpoints **cannot** be in two different child subtrees with LCA `v`, since then `v` (color `x`) would be an internal vertex on the path, violating the rule.
  The only valid paths of color `c[v]` “through `v`” are `(v, w)` where `w` is `c[v]`-top-level in some child subtree. So add:
  `totalCv = ∑_{u child of v} cnt[u][c[v]]`
  to the answer.

#### 4) Small-to-large merge (DSU-on-tree)

To compute `cnt[v]` efficiently:

1. Pick the child `u*` with the largest map `cnt[u*]`. Let `big = cnt[u*]`.
2. Let
   `totalCv = big[c[v]]` (0 if absent)
   then erase key `c[v]` from `big` (descendant top-level `c[v]` must not propagate above `v`).
3. For each other child `u`:
   - For each `(x, val) ∈ cnt[u]`:
     - if `x = c[v]`: `totalCv += val` (do not merge)
     - else: add cross-pairs `ans += big[x] * val`, then `big[x] += val`
4. Add `ans += totalCv` (paths `(v, w)` with endpoint color `c[v]`).
5. Set `big[c[v]] = 1` because `v` itself becomes the only top-level occurrence of its color for ancestors.

Finally set `cnt[v] = big`.

#### 5) Complexity
Small-to-large guarantees each key is merged only `O(log n)` times.
- With ordered maps: `O(n log^2 n)`
- With hash maps (expected): `O(n log n)`

This fits `n ≤ 2·10^5`.