LeetCode 2421 - Number of Good Paths

LeetCode Problem 2421

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Tree, Union-Find, Graph Theory, Sorting

Solution

Problem Understanding

The problem is asking us to count the number of good paths in a tree. A tree is a connected acyclic graph with n nodes and n - 1 edges. Each node has an integer value assigned by the array vals. A good path is defined as a simple path (no repeated nodes) where the first and last node have the same value, and all intermediate nodes have values less than or equal to this maximum value. A single node counts as a valid path because it trivially satisfies the conditions.

The input consists of an array vals of length n, representing node values, and an array edges of length n - 1, representing undirected connections between nodes. The output is a single integer representing the total number of distinct good paths.

The constraints indicate that the number of nodes can be as large as 3 * 10^4, and node values can be as large as 10^5. Because the input is guaranteed to form a tree, we know there are no cycles, every node is connected, and we do not need to check for disconnected components.

Important edge cases include trees with a single node, multiple nodes with the same value, and paths where intermediate nodes exceed the starting node's value. These scenarios can trip up a naive implementation if intermediate values are not carefully considered.

Approaches

Brute Force

A brute-force approach would involve generating all possible simple paths in the tree. For each path, we would check if the start and end nodes have the same value and if all intermediate nodes are less than or equal to this value. While this approach would be correct, it is computationally infeasible. Generating all simple paths in a tree with n nodes can result in exponential complexity, far exceeding the given constraints of n <= 3 * 10^4.

Optimal Approach

The key observation for an optimal solution is that the tree structure allows us to use a Union-Find (Disjoint Set Union) data structure combined with sorting by node values. By processing nodes in increasing order of value, we ensure that when we union nodes, all paths we form satisfy the "all intermediate nodes are less than or equal" constraint. Nodes with the same value can then form good paths with each other. Counting combinations within connected components of nodes with equal values gives us the number of additional good paths beyond the single-node paths.

This approach works because in a tree, any path between two nodes is unique. Processing nodes by value ensures that we only consider valid paths and do not violate the maximum value constraint.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(2^n)	O(n)	Generates all simple paths and validates each one; infeasible for large n
Optimal	O(n log n)	O(n)	Sort nodes by value, use Union-Find to group connected components, count good paths combinatorially

Algorithm Walkthrough

1. Initialize Data Structures: Create a Union-Find structure to manage connected components efficiently. Use a dictionary to map node values to lists of nodes with that value.
2. Sort Nodes by Value: Extract unique values from vals and process them in ascending order. This ensures that when we connect nodes, we respect the "intermediate nodes ≤ start node" constraint.
3. Union Connected Nodes: For each node of the current value, iterate through its neighbors. If a neighbor’s value is less than or equal to the current value, union the two nodes. This step builds connected components that are valid good paths.
4. Count Good Paths in Components: After unioning, for each connected component, count how many nodes have the current value. If there are k such nodes, they form k * (k - 1) / 2 additional good paths (pairwise combinations), plus the k single-node paths.
5. Sum Results: Keep a running total of the number of good paths. Return this value at the end.

Why it works: By processing nodes in order of increasing value and unioning only nodes that respect the value constraint, every good path counted in a component is valid. The tree’s uniqueness of paths guarantees that no path is double-counted, and counting combinations of equal-valued nodes within components accounts for all valid multi-node paths.

Python Solution

from typing import List
from collections import defaultdict

class UnionFind:
    def __init__(self, n: int):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int):
        xr, yr = self.find(x), self.find(y)
        if xr == yr:
            return
        if self.size[xr] < self.size[yr]:
            xr, yr = yr, xr
        self.parent[yr] = xr
        self.size[xr] += self.size[yr]

class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        n = len(vals)
        uf = UnionFind(n)
        graph = [[] for _ in range(n)]
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        val_to_nodes = defaultdict(list)
        for i, val in enumerate(vals):
            val_to_nodes[val].append(i)

        result = 0
        for val in sorted(val_to_nodes):
            nodes = val_to_nodes[val]
            for node in nodes:
                for neighbor in graph[node]:
                    if vals[neighbor] <= val:
                        uf.union(node, neighbor)

            count = defaultdict(int)
            for node in nodes:
                root = uf.find(node)
                count[root] += 1

            for c in count.values():
                result += c * (c + 1) // 2

        return result

Implementation Explanation: We first build the graph adjacency list. Then we map values to the corresponding nodes. Processing values in ascending order allows union operations only among nodes that satisfy the good path constraints. Finally, we count pairwise combinations of nodes in each connected component to determine the total number of good paths.

Go Solution

package main

func numberOfGoodPaths(vals []int, edges [][]int) int {
    n := len(vals)
    parent := make([]int, n)
    size := make([]int, n)
    for i := 0; i < n; i++ {
        parent[i] = i
        size[i] = 1
    }

    var find func(int) int
    find = func(x int) int {
        if parent[x] != x {
            parent[x] = find(parent[x])
        }
        return parent[x]
    }

    union := func(x, y int) {
        xr, yr := find(x), find(y)
        if xr == yr {
            return
        }
        if size[xr] < size[yr] {
            xr, yr = yr, xr
        }
        parent[yr] = xr
        size[xr] += size[yr]
    }

    graph := make([][]int, n)
    for _, e := range edges {
        u, v := e[0], e[1]
        graph[u] = append(graph[u], v)
        graph[v] = append(graph[v], u)
    }

    valToNodes := make(map[int][]int)
    for i, val := range vals {
        valToNodes[val] = append(valToNodes[val], i)
    }

    keys := make([]int, 0, len(valToNodes))
    for k := range valToNodes {
        keys = append(keys, k)
    }
    sort.Ints(keys)

    result := 0
    for _, val := range keys {
        nodes := valToNodes[val]
        for _, node := range nodes {
            for _, neighbor := range graph[node] {
                if vals[neighbor] <= val {
                    union(node, neighbor)
                }
            }
        }

        count := make(map[int]int)
        for _, node := range nodes {
            root := find(node)
            count[root]++
        }

        for _, c := range count {
            result += c * (c + 1) / 2
        }
    }

    return result
}

Go Implementation Notes: Go uses slices for adjacency lists and maps for counting components. The Union-Find is implemented with slices for parent and size. Sorting and map iteration mirrors Python behavior. Integer overflow is not a concern because constraints are within int range.

Worked Examples

Example 1

vals = [1,3,2,1,3]
edges = [[0,1],[0,2],[2,3],[2,4]]

Processing values in order: 1, 2, 3

• Value 1: nodes [0, 3]. Union connected nodes <=1. Components: {0}, {3}. Count: 2 single-node paths + 0 pairs = 2.
• Value 2: node [2]. Union connected nodes <=2. Component: {2}. Count: 1 path.
• Value 3: nodes [1,4]. Union neighbors <=3. 1-0-2-4 connected, root counts: {1:2, 4:1}. Compute pairwise: 3