Graph Weight Queries - Google Top Interview Questions

                        Solution in C++ :

const int MAXN = 100005;
const int MAXD = 17;

// begin union find
int par[MAXN];
int find(int x) {
    return par[x] == x ? x : (par[x] = find(par[x]));
}
bool merge(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y) {
        return false;
    }
    par[x] = y;
    return true;
}
// end union find

// begin structure for MST

vector<pair<int, int>> mstedges[MAXN];  // <destination vertex, weight of edge>
int ancestor[MAXN][MAXD];  // ancestor[i][d] is the node that is 2^d above i in the spanning tree
int ancestorweight[MAXN][MAXD];  // ancestorweight[i][d] is the maximum weight edge over the 2^d
                                 // edges going up from i
int depth[MAXN];                 // depth[i] is the depth of i in its spanning tree

// returns the weight of the heaviest edge of the path from a to b
int getHeaviestEdgeWeight(int a, int b) {
    assert(find(a) == find(b));
    if (depth[a] < depth[b]) {
        swap(a, b);
    }
    int ret = 0;
    for (int d = MAXD - 1; d >= 0; d--) {
        if (depth[a] - depth[b] >= (1 << d)) {
            ret = max(ret, ancestorweight[a][d]);
            a = ancestor[a][d];
        }
    }
    assert(depth[a] == depth[b]);
    for (int d = MAXD - 1; d >= 0; d--) {
        if (ancestor[a][d] != ancestor[b][d]) {
            ret = max(ret, ancestorweight[a][d]);
            ret = max(ret, ancestorweight[b][d]);
            a = ancestor[a][d];
            b = ancestor[b][d];
        }
    }
    if (a != b) {
        ret = max(ret, ancestorweight[a][0]);
        ret = max(ret, ancestorweight[b][0]);
        a = ancestor[a][0];
        b = ancestor[b][0];
    }
    return ret;
}

void initlca(int maxv) {
    for (int d = 1; d < MAXD; d++) {
        for (int i = 0; i <= maxv; i++) {
            ancestor[i][d] = ancestor[ancestor[i][d - 1]][d - 1];
            ancestorweight[i][d] =
                max(ancestorweight[i][d - 1], ancestorweight[ancestor[i][d - 1]][d - 1]);
        }
    }
}

void dfs(int curr, int par) {
    for (auto edge : mstedges[curr]) {
        if (edge.first == par) {
            continue;
        }
        depth[edge.first] = depth[curr] + 1;
        ancestor[edge.first][0] = curr;
        ancestorweight[edge.first][0] = edge.second;
        dfs(edge.first, curr);
    }
}
// end structure for MST

// sorts edges in increasing order of weight
bool edgesort(vector<int>& a, vector<int>& b) {
    return a.back() < b.back();
}

int solve(vector<vector<int>>& edges, vector<vector<int>>& queries) {
    sort(edges.begin(), edges.end(), edgesort);
    int maxv = 0;
    for (auto& out : edges) {
        maxv = max(maxv, max(out[0], out[1]));
    }
    for (int i = 0; i <= maxv; i++) {
        par[i] = i;
    }
    for (int i = 0; i <= maxv; i++) {
        mstedges[i].clear();
    }
    for (auto& edge : edges) {
        if (merge(edge[0], edge[1])) {
            mstedges[edge[0]].emplace_back(edge[1], edge[2]);
            mstedges[edge[1]].emplace_back(edge[0], edge[2]);
        }
    }
    for (int i = 0; i <= maxv; i++) {
        depth[i] = -1;
    }
    for (int i = 0; i <= maxv; i++) {
        if (depth[i] == -1) {
            depth[i] = 0;
            ancestor[i][0] = i;
            ancestor[i][0] = 0;
            dfs(i, -1);
        }
    }
    initlca(maxv);
    int ret = 0;
    for (auto& query : queries) {
        ret += find(query[0]) == find(query[1]) &&
               query[2] >= getHeaviestEdgeWeight(query[0], query[1]);
    }
    return ret;
}
                    

                        Solution in Java :

import java.util.*;

/**
 * This is a implementation class of union by rank(union by weight) with path compression.
 */
class DisjointSetUnion {
    List<Integer> disJointSet;

    public DisjointSetUnion(int n) {
        disJointSet = new ArrayList<Integer>();
        for (int i = 0; i < n; ++i) {
            disJointSet.add(-1);
        }
    }

    public void union(int firstNode, int secondNode) {
        int firstNodeParent = find(firstNode);
        int secondNodeParent = find(secondNode);
        if (firstNodeParent == secondNodeParent)
            return;
        if (Math.abs(disJointSet.get(firstNodeParent))
            >= Math.abs(disJointSet.get(secondNodeParent))) {
            disJointSet.set(firstNodeParent,
                disJointSet.get(firstNodeParent) + disJointSet.get(secondNodeParent));
            disJointSet.set(secondNodeParent, firstNodeParent);
        } else {
            disJointSet.set(secondNodeParent,
                disJointSet.get(secondNodeParent) + disJointSet.get(firstNodeParent));
            disJointSet.set(firstNodeParent, secondNodeParent);
        }
    }

    public int find(int node) {
        if (disJointSet.get(node) < 0) {
            return node;
        }
        int parent = find(disJointSet.get(node));
        disJointSet.set(node, parent);
        return parent;
    }

    public void printDisjointSet() {
        disJointSet.forEach(element -> System.out.print(element + " "));
    }
}

class Edge implements Comparable {
    private int u, v, w;

    public Edge(int u, int v, int w) {
        this.u = u;
        this.v = v;
        this.w = w;
    }

    public int getU() {
        return u;
    }

    public int getV() {
        return v;
    }

    public int getW() {
        return w;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o)
            return true;
        if (o == null || getClass() != o.getClass())
            return false;
        Edge edge = (Edge) o;
        return u == edge.u && v == edge.v && w == edge.w;
    }

    @Override
    public int hashCode() {
        return Objects.hash(u, v, w);
    }

    @Override
    public int compareTo(Object o) {
        Edge node = (Edge) o;
        return this.w - node.getW();
    }
}

class Solution {
    public int solve(int[][] e, int[][] query) {
        int noOfVertices = Integer.MIN_VALUE;
        int n = e.length;
        int q = query.length;
        List<Edge> edges = new ArrayList<>();
        for (int i = 0; i < n; ++i) {
            int u = e[i][0];
            int v = e[i][1];
            int w = e[i][2];
            noOfVertices = Math.max(noOfVertices, Math.max(u, v));
            Edge edge = new Edge(u, v, w);
            edges.add(edge);
        }
        List<Edge> queries = new ArrayList<>();
        for (int i = 0; i < q; ++i) {
            int u = query[i][0];
            int v = query[i][1];
            int w = query[i][2];
            Edge edge = new Edge(u, v, w);
            queries.add(edge);
        }
        Collections.sort(queries);
        Collections.sort(edges);
        int edgesPosition = 0;
        int queriesPosition = 0;
        DisjointSetUnion disjointSetUnion = new DisjointSetUnion(noOfVertices + 1);
        int answer = 0;
        while (queriesPosition < queries.size()) {
            int val = queries.get(queriesPosition).getW();
            edgesPosition = eddEdgesUptillValue(disjointSetUnion, edges, edgesPosition, val);
            if (disjointSetUnion.find(queries.get(queriesPosition).getU())
                == disjointSetUnion.find(queries.get(queriesPosition).getV())) {
                ++answer;
            }
            ++queriesPosition;
        }
        return answer;
    }
    private int eddEdgesUptillValue(
        DisjointSetUnion disjointSetUnion, List<Edge> edges, int edgesPosition, int val) {
        while (edgesPosition < edges.size() && edges.get(edgesPosition).getW() <= val) {
            disjointSetUnion.union(
                edges.get(edgesPosition).getU(), edges.get(edgesPosition).getV());
            ++edgesPosition;
        }
        return edgesPosition;
    }
}
                    

                        Solution in Python : 
                            
class DSU:
    def __init__(self, N):
        self.par = list(range(N))
        self.sz = [1] * N

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

    def union(self, x, y):
        xr, yr = self.find(x), self.find(y)
        if xr == yr:
            return False
        if self.sz[xr] < self.sz[yr]:
            xr, yr = yr, xr
        self.par[yr] = xr
        self.sz[xr] += self.sz[yr]
        self.sz[yr] = self.sz[xr]
        return True


class Solution:
    def solve(self, edges, queries):
        N = max(max(u, v) for u, v, w in edges) + 1
        edges.sort(key=itemgetter(2))
        queries.sort(key=itemgetter(2))

        dsu = DSU(N)
        ans = i = 0
        for u, v, w in queries:
            while i < len(edges) and edges[i][2] <= w:
                dsu.union(*edges[i][:2])
                i += 1
            if dsu.find(u) == dsu.find(v):
                ans += 1

        return ans