Problem statement

https://binarysearch.com/problems/Graph-Weight-Queries/

Solution

Variation of Leetcode 1697. Checking Existence of Edge Length Limited Paths, but first we need to use <=, not < for weight and also add n as parameter of problem

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 solve(self, edgeList, queries):
        n = 0
        for x, y, _ in edgeList:
            n = max(n, x + 1, y + 1)

        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 sum(ans)