Checking Existence of Edge Length Limited Paths

hard graph union find sorting

Problem

An undirected graph on n nodes has weighted edges [u, v, dist]. For each query [p, q, limit], answer whether there is a path from p to q using only edges of distance strictly less than limit. Return a boolean array of answers.

Sort edges by weight and queries by limit. Process queries in increasing limit; before each, union every edge whose weight is below that limit. Then the query is true iff p and q share a root.

Inputn = 3, edges = [[0,1,2],[1,2,4],[2,0,8],[1,0,16]], queries = [[0,1,2],[0,2,5]]
Output[false, true]
For limit 2 no edge < 2 exists, so 0–1 is false. For limit 5 edges 2 and 4 connect 0–1–2, so 0–2 is true.

def distance_limited_paths(n, edge_list, queries):
    parent = list(range(n))
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    def union(a, b):
        parent[find(a)] = find(b)
    edge_list.sort(key=lambda e: e[2])
    order = sorted(range(len(queries)), key=lambda i: queries[i][2])
    ans = [False] * len(queries)
    i = 0
    for qi in order:
        p, q, limit = queries[qi]
        while i < len(edge_list) and edge_list[i][2] < limit:
            union(edge_list[i][0], edge_list[i][1])
            i += 1
        ans[qi] = find(p) == find(q)
    return ans
function distanceLimitedPaths(n, edgeList, queries) {
  const parent = Array.from({ length: n }, (_, i) => i);
  function find(x) {
    while (parent[x] !== x) { parent[x] = parent[parent[x]]; x = parent[x]; }
    return x;
  }
  function union(a, b) { parent[find(a)] = find(b); }
  edgeList.sort((a, b) => a[2] - b[2]);
  const order = queries.map((_, i) => i).sort((a, b) => queries[a][2] - queries[b][2]);
  const ans = new Array(queries.length).fill(false);
  let i = 0;
  for (const qi of order) {
    const [p, q, limit] = queries[qi];
    while (i < edgeList.length && edgeList[i][2] < limit) {
      union(edgeList[i][0], edgeList[i][1]);
      i++;
    }
    ans[qi] = find(p) === find(q);
  }
  return ans;
}
class Solution {
    int[] parent;
    int find(int x) {
        while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; }
        return x;
    }
    void union(int a, int b) { parent[find(a)] = find(b); }
    public boolean[] distanceLimitedPaths(int n, int[][] edgeList, int[][] queries) {
        parent = new int[n];
        for (int i = 0; i < n; i++) parent[i] = i;
        Arrays.sort(edgeList, (a, b) -> a[2] - b[2]);
        Integer[] order = new Integer[queries.length];
        for (int i = 0; i < order.length; i++) order[i] = i;
        Arrays.sort(order, (a, b) -> queries[a][2] - queries[b][2]);
        boolean[] ans = new boolean[queries.length];
        int i = 0;
        for (int qi : order) {
            int p = queries[qi][0], q = queries[qi][1], limit = queries[qi][2];
            while (i < edgeList.length && edgeList[i][2] < limit) {
                union(edgeList[i][0], edgeList[i][1]);
                i++;
            }
            ans[qi] = find(p) == find(q);
        }
        return ans;
    }
}
class Solution {
public:
    vector<int> parent;
    int find(int x) {
        while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; }
        return x;
    }
    void uni(int a, int b) { parent[find(a)] = find(b); }
    vector<bool> distanceLimitedPaths(int n, vector<vector<int>>& edgeList, vector<vector<int>>& queries) {
        parent.resize(n);
        for (int i = 0; i < n; i++) parent[i] = i;
        sort(edgeList.begin(), edgeList.end(), [](auto& a, auto& b){ return a[2] < b[2]; });
        vector<int> order(queries.size());
        for (int i = 0; i < (int)order.size(); i++) order[i] = i;
        sort(order.begin(), order.end(), [&](int a, int b){ return queries[a][2] < queries[b][2]; });
        vector<bool> ans(queries.size(), false);
        int i = 0;
        for (int qi : order) {
            int p = queries[qi][0], q = queries[qi][1], limit = queries[qi][2];
            while (i < (int)edgeList.size() && edgeList[i][2] < limit) {
                uni(edgeList[i][0], edgeList[i][1]);
                i++;
            }
            ans[qi] = find(p) == find(q);
        }
        return ans;
    }
};
Time: O(E log E + Q log Q) Space: O(n + Q)