Intuition

  1. All single node are considered as a good path.
  2. Take a look at the biggest value maxv
    If the number of nodes with biggest value maxv is k,
    then good path starting with value maxv is k*(k-1)/2,
    any pair can build up a good path.
    And these nodes will cut the tree into small trees.

Explanation

We do the "cutting" process reversely by union-find.

Firstly we initilize result res = n,
we union the edge with small values v.

Then the nodes of value v on one side,
can pair up with the nodes of value v on the other side.

We union one node into the other by union action,
and also the count of nodes with value v.

We reapeatly doing union action for all edges until we build up the whole tree.

Complexity

Time O(union-find(n))
Space O(union-find(n))

Python

    def numberOfGoodPaths(self, vals, edges):
        res = n = len(vals)
        f = list(range(n))
        count = [Counter({vals[i]: 1}) for i in range(n)]
        edges = sorted([max(vals[i], vals[j]),i,j] for i,j in edges)

        def find(x):
            if f[x] != x:
                f[x] = find(f[x])
            return f[x]

        for v,i,j in edges:
            fi, fj = find(i), find(j)
            cj, ci = count[fi][v], count[fj][v]
            res += ci * cj
            f[fj] = fi
            count[fi] = Counter({v: ci + cj})
        return res