Intuition
- All single node are considered as a good path.
- Take a look at the biggest value
maxv
If the number of nodes with biggest valuemaxv
isk
,
then good path starting with valuemaxv
isk*(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