<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>

from typing import List

class UnionFind:
    def __init__(self, n, vals):
        self.parent = list(range(n))  # Each node is initially its own parent.
        self.size = [1] * n  # Initialize size for each component to 1.
        self.max_val_seen = vals[:]  # Initialize max value seen in each component.
        self.count = [1] * n  # Initialize count of max value nodes in each component.

    def find(self, node):
        if node == self.parent[node]:
            return node
        # Path compression: make the parent of 'node' point to the root directly
        self.parent[node] = self.find(self.parent[node])
        return self.parent[node]

    def union(self, node1, node2):
        root1, root2 = self.find(node1), self.find(node2)

        if root1 != root2:
            # Ensure root1 is always the larger root to maintain a smaller tree height
            if self.size[root1] < self.size[root2]:
                root1, root2 = root2, root1
            # Union operation: point the root of the smaller tree to the root of the larger tree
            self.parent[root2] = root1
            # Update the size of the new root
            self.size[root1] += self.size[root2]

            # If the maximum values in both components are the same, merge and count good paths
            if self.max_val_seen[root1] == self.max_val_seen[root2]:
                # Example:
                # root1's component: max_val_seen = 3, count = 2 (nodes: 1, 4)
                # root2's component: max_val_seen = 3, count = 1 (nodes: 2)
                # New good paths formed by merging:
                # Good paths = count[root1] * count[root2] = 2 * 1 = 2

                # The number of new good paths is the product of the counts of max value nodes in both components.
                result = self.count[root1] * self.count[root2]  # New good paths

                # Update the count of max value nodes in the new root component.
                # Explanation:
                # Updated count = count[root1] + count[root2] = 2 + 1 = 3
                self.count[root1] += self.count[root2]
                return result

            # If the maximum value in root1's component is less than root2's component, update root1's component
            elif self.max_val_seen[root1] < self.max_val_seen[root2]:
                # Example:
                # root1's component: max_val_seen = 2, count = 1 (nodes: 0)
                # root2's component: max_val_seen = 3, count = 2 (nodes: 1, 4)
                # After merging, root1's component should reflect the max value and count of root2's.

                # Set the max value seen in root1's component to root2's max value.
                self.max_val_seen[root1] = self.max_val_seen[root2]

                # Update the count of max value nodes in root1's component to root2's count.
                # Explanation:
                # Updated count = count[root2] = 2
                self.count[root1] = self.count[root2]
        return 0

class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        """
        This function counts all the good paths as defined in the problem statement
        by gradually building the graph using a union-find structure to track connected
        components efficiently.
        """
        n = len(vals)
        ans = n  # Start with each node as a single-node good path
        uf = UnionFind(n, vals)
        
        # Sort edges based on the maximum value of the connected nodes to ensure the correct order of processing.
        edges.sort(key=lambda edge: max(vals[edge[0]], vals[edge[1]]))

        # Process each edge, connecting nodes and counting additional good paths formed
        for a, b in edges:
            ans += uf.union(a, b)
            
        return ans

# Test the solution with specific cases
vals1 = [1, 3, 2, 1, 3]
edges1 = [[0, 1], [0, 2], [2, 3], [2, 4]]
print(Solution().numberOfGoodPaths(vals1, edges1))  # Expected: 6

vals2 = [1, 1, 2, 2, 3]
edges2 = [[0, 1], [1, 2], [2, 3], [2, 4]]
print(Solution().numberOfGoodPaths(vals2, edges2))  # Expected: 7

vals3 = [1]
edges3 = []
print(Solution().numberOfGoodPaths(vals3, edges3))  # Expected: 1

# Explanation and examples:
"""
Example:
vals = [1, 3, 2, 1, 3]
edges = [[0, 1], [0, 2], [2, 3], [2, 4]]

Initial state:
Nodes:      0 1 2 3 4
Vals:       1 3 2 1 3
Parent:     0 1 2 3 4
Size:       1 1 1 1 1
MaxValSeen: 1 3 2 1 3
Count:      1 1 1 1 1

# After connecting nodes 2 and 4 (same max values):
# Node 2 with value 2, Node 4 with value 3
# Root1 = 2, Root2 = 4
# MaxValSeen[root1] < MaxValSeen[root2], so update root1's component.
Parent:     0 1 4 3 4  # Node 2 points to node 4
Size:       1 1 2 1 2  # Size of the tree rooted at 4 is now 2
MaxValSeen: 1 3 3 1 3  # Max values remain unchanged
Count:      1 1 1 1 1  # Count remains unchanged

# After connecting nodes 0 and 1 (max values in the same components):
# Node 0 with value 1, Node 1 with value 3
# Root1 = 0, Root2 = 1
# MaxValSeen[root1] < MaxValSeen[root2], so update root1's component.
Parent:     0 0 4 3 4  # Node 1 points to node 0, merged under root 0
Size:       2 1 2 1 2  # Size of the tree rooted at 0 is now 2
MaxValSeen: 3 3 3 1 3  # Max value of component rooted at 0 updated to 3
Count:      1 1 1 1 1  # Count of max value nodes in component rooted at 0 updated to 1
Result:     0  # No new good paths as max values are different

# After connecting nodes 2 and 3 (max values in the same components):
# Node 2 with value 3, Node 3 with value 1
# Root1 = 4, Root2 = 3
# MaxValSeen[root1] > MaxValSeen[root2], no change in max value
Parent:     0 0 4 4 4  # Node 3 points to node 4, merged under root 4
Size:       2 1 3 1 4  # Size of the tree rooted at 4 is now 4
MaxValSeen: 3 3 3 3 3  # Max values remain unchanged
Count:      1 1 1 1 1  # Count remains unchanged
Result:     0  # No new good paths as max values are different

# After connecting nodes 0 and 2 (max values in the same components):
# Node 0 with value 3, Node 2 with value 3
# Root1 = 0, Root2 = 4
# MaxValSeen[root1] == MaxValSeen[root2], calculate new good paths
Parent:     0 0 4 4 0  # Node 4 points to node 0, merged under root 0
Size:       4 1 3 1 4  # Size of the tree rooted at 0 is now 4
MaxValSeen: 3 3 3 3 3  # Max values remain unchanged
Count:      3 1 1 1 1  # Count of max value nodes in component rooted at 0 updated to 3
Result:     1  # 1 good path from merging [0, 2] where max values are equal

Diagram before and after connecting:
Before:
       0 (1)
      /
     1 (3)
     
       2 (2)
      /
     4 (3)
     
After connecting 0 and 2:
       0 (3)
      /|\
     1 2 4
"""

1181