Problem statement

https://leetcode.com/problems/checking-existence-of-edge-length-limited-paths/

Solution

The idea is to sort edges by weight and add them one by one, using union find to merge connected components.

Complexity

It is O(n log n + q log q) for time and O(n + q) for space.

Code

class DSU:
    def __init__(self, N):
        self.p = list(range(N))

    def find(self, x):
        if self.p[x] != x:
            self.p[x] = self.find(self.p[x])
        return self.p[x]

    def union(self, x, y):
        xr = self.find(x)
        yr = self.find(y)
        self.p[xr] = yr

class Solution:
    def distanceLimitedPathsExist(self, n, edgeList, queries):
        dsu = DSU(n)
        m, L, e = len(queries), len(edgeList), 0
        
        q_sort = sorted((w, x, y, i) for i, (x, y, w) in enumerate(queries))
        e_sort = sorted(edgeList, key = lambda x: x[2])
        ans = [0] * len(queries)
        
        for q in range(m):
            while e < L and e_sort[e][2] < q_sort[q][0]:
                dsu.union(e_sort[e][0], e_sort[e][1])
                e += 1
            ans[q_sort[q][3]] = dsu.find(q_sort[q][1]) == dsu.find(q_sort[q][2])
            
        return ans