main.cpp

// Source: https://leetcode.com/problems/checking-existence-of-edge-length-limited-paths-ii
// Title: Checking Existence of Edge Length Limited Paths II
// Difficulty: Hard
// Author: Mu Yang <http://muyang.pro>

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// An undirected graph of `n` nodes is defined by `edgeList`, where `edgeList[i] = [u_i, v_i, dis_i]` denotes an edge between nodes `u_i` and `v_i` with distance `dis_i`. Note that there may be **multiple** edges between two nodes, and the graph may not be connected.
//
// Implement the `DistanceLimitedPathsExist` class:
//
// - `DistanceLimitedPathsExist(int n, int[][] edgeList)` Initializes the class with an undirected graph.
// - `boolean query(int p, int q, int limit)` Returns `true` if there exists a path from `p` to `q` such that each edge on the path has a distance **strictly less than** `limit`, and otherwise `false`.
//
// **Example 1:**
//
// **https://assets.leetcode.com/uploads/2021/01/05/messed.png**
//
// ```
// Input
//
// ["DistanceLimitedPathsExist", "query", "query", "query", "query"]
// [[6, [[0, 2, 4], [0, 3, 2], [1, 2, 3], [2, 3, 1], [4, 5, 5]]], [2, 3, 2], [1, 3, 3], [2, 0, 3], [0, 5, 6]]
// Output
//
// [null, true, false, true, false]
//
// Explanation
//
// DistanceLimitedPathsExist distanceLimitedPathsExist = new DistanceLimitedPathsExist(6, [[0, 2, 4], [0, 3, 2], [1, 2, 3], [2, 3, 1], [4, 5, 5]]);
// distanceLimitedPathsExist.query(2, 3, 2); // return true. There is an edge from 2 to 3 of distance 1, which is less than 2.
// distanceLimitedPathsExist.query(1, 3, 3); // return false. There is no way to go from 1 to 3 with distances **strictly** less than 3.
// distanceLimitedPathsExist.query(2, 0, 3); // return true. There is a way to go from 2 to 0 with distance < 3: travel from 2 to 3 to 0.
// distanceLimitedPathsExist.query(0, 5, 6); // return false. There are no paths from 0 to 5.
// ```
//
// `**Constraints:**`
//
// - `2 <= n <= 10^4`
// - `0 <= edgeList.length <= 10^4`
// - `edgeList[i].length == 3`
// - `0 <= u_i, v_i, p, q <= n-1`
// - `u_i != v_i`
// - `p != q`
// - `1 <= dis_i, limit <= 10^9`
// - At most`10^4` calls will be made to `query`.
//
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include <climits>
#include <numeric>
#include <vector>

using namespace std;

// Union-Find + Sort
//
// First the edges by its weight.
// Next use union find to put the edges into the graph one by one.
//
// In this union find, we skip path compression in find.
// We also store the edge weight in the union find.
//
// For query, we find the traverse up in union find for paths smaller than limit.
class DistanceLimitedPathsExist {
  struct UnionFind {
    vector<int> parents;
    vector<int> ranks;
    vector<int> weights;  // weight of edge to parent

    UnionFind(int n) : parents(n), ranks(n), weights(n, INT_MAX) {  //
      iota(parents.begin(), parents.end(), 0);
    }

    // Find with weight < limit
    int find(int x, int limit) const {
      if (weights[x] >= limit) return x;
      return find(parents[x], limit);
    }

    bool isConnected(int x, int y, int limit) const {  //
      return find(x, limit) == find(y, limit);
    }

    void unite(int x, int y, int weight) {
      x = find(x, INT_MAX);
      y = find(y, INT_MAX);
      if (x == y) return;

      // Ensure rank(x) >= rank(y)
      if (ranks[x] < ranks[y]) swap(x, y);

      // Merge y into x
      if (ranks[x] == ranks[y]) ++ranks[x];
      parents[y] = x;
      weights[y] = weight;
    }
  };

  vector<vector<int>>& edges;
  UnionFind uf;

 public:
  DistanceLimitedPathsExist(int n, vector<vector<int>>& edges) : edges(edges), uf(n) {
    sort(edges.begin(), edges.end(), [](const vector<int>& a, const vector<int>& b) {  //
      return a[2] < b[2];
    });

    for (const auto& edge : edges) {
      int u = edge[0], v = edge[1], w = edge[2];
      uf.unite(u, v, w);
    }
  }

  bool query(int p, int q, int limit) {  //
    return uf.isConnected(p, q, limit);
  }
};