Checking Existence of Edge Length Limited Paths

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.

Given an array queries, where queries[j] = [pj, qj, limitj], your task is to determine for each queries[j] whether there is a path between pj and qj such that each edge on the path has a distance strictly less than limitj .

Return a boolean array answer, where answer.length == queries.length and the jth value of answer is true if there is a path for queries[j] is true, and false otherwise.

 

Example 1:

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.
For the first query, between 0 and 1 there is no path where each distance is less than 2, thus we return false for this query.
For the second query, there is a path (0 -> 1 -> 2) of two edges with distances less than 5, thus we return true for this query.

Example 2:

Input: n = 5, edgeList = [[0,1,10],[1,2,5],[2,3,9],[3,4,13]], queries = [[0,4,14],[1,4,13]]
Output: [true,false]
Exaplanation: The above figure shows the given graph.

 

Constraints:


Solution:

class Solution {
    static class UF {
        int[] parent;
        int[] size;
        
        public UF(int n) {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i ++) {
                parent[i] = i;
                size[i] = 1;
            }
        }
        
        public void union(int x, int y) {
            int rootX = find(x);
            int rootY = find(y);
            if (rootX == rootY) {
                return;
            }
            if (size[rootX] > size[rootY]) {
                parent[rootY] = rootX;
                size[rootX] += size[rootY];
            } else {
                parent[rootX] = rootY;
                size[rootY] += size[rootX];
            }
        }
        
        public int find(int x) {
            while (x != parent[x]) {
                parent[x] = parent[parent[x]];
                x = parent[x];
            }
            return x;
        }
        
        public boolean connected(int x, int y) {
            return find(x) == find(y);
        }
    }
    
    
    public boolean[] distanceLimitedPathsExist(int n, int[][] edgeList, int[][] queries) {
        int q = queries.length;
        int e = edgeList.length;
        boolean[] res = new boolean[q];
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>((a, b) -> a[2] - b[2]);
        for (int i = 0; i < q; i ++) {
            int[] query = queries[i];
            int[] qq = new int[]{query[0], query[1], query[2], i};
            pq.offer(qq);
        }
        Arrays.sort(edgeList, (a, b) -> a[2] - b[2]);
        UF nodes = new UF(n);
        int i = 0;
        while (!pq.isEmpty()) {
            int[] curr = pq.poll();
            int currMax = curr[2];
            while (i < e && edgeList[i][2] < currMax) {
                int[] edge = edgeList[i];
                nodes.union(edge[0], edge[1]);
                i ++;
            }
            if (nodes.connected(curr[0], curr[1])) {
                res[curr[3]] = true;
            }
        }
        return res;
    }
}