<p>There is a tree (i.e. a connected, undirected graph with no cycles) consisting of <code>n</code> nodes numbered from <code>0</code> to <code>n - 1</code> and exactly <code>n - 1</code> edges.</p>

<p>You are given a <strong>0-indexed</strong> integer array <code>vals</code> of length <code>n</code> where <code>vals[i]</code> denotes the value of the <code>i<sup>th</sup></code> node. You are also given a 2D integer array <code>edges</code> where <code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code> denotes that there exists an <strong>undirected</strong> edge connecting nodes <code>a<sub>i</sub></code> and <code>b<sub>i</sub></code>.</p>

<p>A <strong>good path</strong> is a simple path that satisfies the following conditions:</p>

<ol>
	<li>The starting node and the ending node have the <strong>same</strong> value.</li>
	<li>All nodes between the starting node and the ending node have values <strong>less than or equal to</strong> the starting node (i.e. the starting node&#39;s value should be the maximum value along the path).</li>
</ol>

<p>Return <em>the number of distinct good paths</em>.</p>

<p>Note that a path and its reverse are counted as the <strong>same</strong> path. For example, <code>0 -&gt; 1</code> is considered to be the same as <code>1 -&gt; 0</code>. A single node is also considered as a valid path.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2022/08/04/f9caaac15b383af9115c5586779dec5.png" style="width: 400px; height: 333px;" />
<pre>
<strong>Input:</strong> vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]
<strong>Output:</strong> 6
<strong>Explanation:</strong> There are 5 good paths consisting of a single node.
There is 1 additional good path: 1 -&gt; 0 -&gt; 2 -&gt; 4.
(The reverse path 4 -&gt; 2 -&gt; 0 -&gt; 1 is treated as the same as 1 -&gt; 0 -&gt; 2 -&gt; 4.)
Note that 0 -&gt; 2 -&gt; 3 is not a good path because vals[2] &gt; vals[0].
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2022/08/04/149d3065ec165a71a1b9aec890776ff.png" style="width: 273px; height: 350px;" />
<pre>
<strong>Input:</strong> vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]
<strong>Output:</strong> 7
<strong>Explanation:</strong> There are 5 good paths consisting of a single node.
There are 2 additional good paths: 0 -&gt; 1 and 2 -&gt; 3.
</pre>

<p><strong class="example">Example 3:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2022/08/04/31705e22af3d9c0a557459bc7d1b62d.png" style="width: 100px; height: 88px;" />
<pre>
<strong>Input:</strong> vals = [1], edges = []
<strong>Output:</strong> 1
<strong>Explanation:</strong> The tree consists of only one node, so there is one good path.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == vals.length</code></li>
	<li><code>1 &lt;= n &lt;= 3 * 10<sup>4</sup></code></li>
	<li><code>0 &lt;= vals[i] &lt;= 10<sup>5</sup></code></li>
	<li><code>edges.length == n - 1</code></li>
	<li><code>edges[i].length == 2</code></li>
	<li><code>0 &lt;= a<sub>i</sub>, b<sub>i</sub> &lt; n</code></li>
	<li><code>a<sub>i</sub> != b<sub>i</sub></code></li>
	<li><code>edges</code> represents a valid tree.</li>
</ul>

"""
Approach: Union Find + Counting + Sorting
    0. We view each edge as the bridge to introduce new good paths:
        Edge (u, v):
            val = max(u.val, v.val)
            node1 = u
            node2 = v
        1. When a new edge is introduced, that means two disconnected groups will be connected, 
        and we can find new good paths by finding one endpoint from each group

    1. Sort the edge by their value:
        1. This ensures that we only consider the current maximum value to build the good paths
        2. For example, this can avoid adding new good paths of smaller value
                     u
        group 1: 1 - 2
                     |
        group 2:     1
                     v
                    In this case, edge (u, v) does not increase the # good paths, as the node in between (2) 
                    is larger than two endpoints (1) 


    2. We traverse each edge to connect two groups:
        1. We use edge.val as the current maximum value
        2. When two groups are connected, that means num_good_paths += group1[edge.val] * group2[edge2.val]
        For instance, say we have these two groups:
                            u
            group1: 0 - 2 - 2
                            | (new edge)
            group2:     1 - 2
                            v
        
        If we connect two groups with edge (u, v):
            0. the value of the edge would be 2, we use this value as the value of good path,
                that is, we want to find one endpoint from group1 and another endpoint from group2, 
                both with value 2
            1. We will find 2 * 1 = 2 new good paths, as there are two nodes with value 2 in group 1,
                and one node with value 2 in group 2
        3. To do this, we need to count the frequency of the value for each group:
            root_counter[root][val] = freq
        4. We use Union-Find to maintain the group info and to handle the group merge
    3. Finally num_good_paths should be further increased by # nodes
    4. return num_good_paths

TC: O(ElgE + E + V), E denotes # edges, V denotes # nodes
    Sorting the edges - O(ElgE)
    Using Disjoint Set for Union Find takes O(1) on average, and we need O(E + V) operations

SC: O(E + V)
    O(E) to store edges
    O(V) for the counter (group1[v])
"""
class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        def find_root(node):
            if root[node] != node:
                root[node] = find_root(root[node])
            return root[node]
        
        def union(val, node1, node2):
            node1 = find_root(node1)
            node2 = find_root(node2)

            if node1 != node2:
                if size[node1] > size[node2]:
                    node1, node2 = node2, node1
                
                root[node1] = node2
                size[node2] += size[node1]
                root_counter[node2][val] += root_counter[node1][val]
            return (root_counter[node2][val] - root_counter[node1][val]) * root_counter[node1][val]

        # Build and sort the edges by value.
        edges = sorted([(max(vals[node1], vals[node2]), node1, node2) for node1, node2 in edges])

        # Initialization for disjoint set.
        n = len(vals)
        root = [i for i in range(n)]
        size = [1 for _ in range(n)]
        root_counter = [Counter({vals[i]: 1}) for i in range(n)]

        # Run Union Find to merge groups
        num_good_paths = 0
        for edge in edges:
            num_good_paths += union(*edge)
        return num_good_paths + n



        


        
1897