Number of Good Paths Python Solution

Tested Python solution for LeetCode 2421 with 23 pytest cases. Generate a practice environment with lcpy.

LeetCode 2421, Hard. Topics: Array, Hash Table, Tree, Union-Find, Graph Theory, Sorting. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 23 parametrized pytest cases, and a playground notebook:

lcpy gen -n 2421   # by problem number
lcpy gen -s number_of_good_paths   # by problem name

​

Problem

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

You are given a 0-indexed integer array vals of length n where vals[i] denotes the value of the ith node. You are also given a 2D integer array edges where edges[i] = [ai, bi] denotes that there exists an undirected edge connecting nodes ai and bi.

A good path is a simple path that satisfies the following conditions:

1. The starting node and the ending node have the same value.
2. All nodes between the starting node and the ending node have values less than or equal to the starting node (i.e. the starting node’s value should be the maximum value along the path).

Return the number of distinct good paths.

Note that a path and its reverse are counted as the same path. For example, 0 -> 1 is considered to be the same as 1 -> 0. A single node is also considered as a valid path.

​

Examples

Input: vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]
Output: 6
Explanation: There are 5 good paths consisting of a single node.
There is 1 additional good path: 1 -> 0 -> 2 -> 4.
(The reverse path 4 -> 2 -> 0 -> 1 is treated as the same as 1 -> 0 -> 2 -> 4.)
Note that 0 -> 2 -> 3 is not a good path because vals[2] > vals[0].

Input: vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]
Output: 7
Explanation: There are 5 good paths consisting of a single node.
There are 2 additional good paths: 0 -> 1 and 2 -> 3.

Input: vals = [1], edges = []
Output: 1
Explanation: The tree consists of only one node, so there is one good path.

​

Constraints

• n == vals.length
• 1 <= n <= 3 * 10^4
• 0 <= vals[i] <= 10^5
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• edges represents a valid tree.

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

class Solution:
    # Time: O(n log n + e * alpha(n)) for sorting the nodes and unioning edges
    # Space: O(n) for the parent, size, adjacency and active arrays
    def number_of_good_paths(self, vals: list[int], edges: list[list[int]]) -> int:
        n = len(vals)
        parent = list(range(n))
        size = [1] * n
        active = [False] * n

        adj: list[list[int]] = [[] for _ in range(n)]
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)

        def find(x: int) -> int:
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        def union(a: int, b: int) -> None:
            ra, rb = find(a), find(b)
            if ra == rb:
                return
            if size[ra] < size[rb]:
                ra, rb = rb, ra
            parent[rb] = ra
            size[ra] += size[rb]

        # Every single node is a good path on its own.
        result = n

        # Grow the forest from low values to high values: paths whose maximum
        # value is v are only complete once every node with value <= v exists.
        order = sorted(range(n), key=lambda i: vals[i])
        i = 0
        while i < n:
            j = i
            while j < n and vals[order[j]] == vals[order[i]]:
                j += 1
            group = order[i:j]

            for node in group:
                active[node] = True
            for node in group:
                for nxt in adj[node]:
                    if active[nxt]:
                        union(node, nxt)

            # Every pair of value-v nodes sharing a component gives one path.
            counts: dict[int, int] = {}
            for node in group:
                root = find(node)
                counts[root] = counts.get(root, 0) + 1
            for c in counts.values():
                result += c * (c - 1) // 2

            i = j
        return result

​

Complexity

Time	Space
O(n log n + e * alpha(n)) for sorting the nodes and unioning edges	O(n) for the parent, size, adjacency and active arrays

​

Tags

NeetCode All.