Checking Existence of Edge Length Limited Paths — LeetCode 1697 Python Solution

The problem

An undirected graph of n nodes is defined by edgeList, where edgeList[i] = [ui, vi, disi] denotes an edge between nodes ui and vi with distance disi. Note that there may be multiple edges between two nodes.

Example

Input

n = 3, edgeList = [[0,1,2],[1,2,4],[2,0,8],[1,0,16]], queries = [[0,1,2],[0,2,5]]

Output

[false,true]

Explanation

The above figure shows the given graph. Note that there are two overlapping edges between 0 and 1 with distances 2 and 16.

Python solution

class Solution:
    def distanceLimitedPathsExist(
        self, n: int, edgeList: List[List[int]], queries: List[List[int]]
    ) -> List[bool]:
        def find(x):
            if p[x] != x:
                p[x] = find(p[x])
            return p[x]

        p = list(range(n))
        edgeList.sort(key=lambda x: x[2])
        j = 0
        ans = [False] * len(queries)
        for i, (a, b, limit) in sorted(enumerate(queries), key=lambda x: x[1][2]):
            while j < len(edgeList) and edgeList[j][2] < limit:
                u, v, _ = edgeList[j]
                p[find(u)] = find(v)
                j += 1
            ans[i] = find(a) == find(b)
        return ans

Complexity

Pattern: Union-Find

Merge groups and ask whether two things are connected, both in near-constant time. LeetCode 1697. Checking Existence of Edge Length Limited Paths is filed here on both counts: the reference solution below belongs to the algorithm family this hub collects, and LeetCode tags it Union Find.

The union-find guide has the Python template for the pattern and the 83 LeetCode problems that use it.

Related problems