Skip to content

3534. Path Existence Queries in a Graph II 👍

  • Time: $O(q\log n)$
  • Space: $O(n\log n + q)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class Solution {
 public:
  vector<int> pathExistenceQueries(int n, vector<int>& nums, int maxDiff,
                                   vector<vector<int>>& queries) {
    vector<int> ans;
    vector<int> sortedNums;
    vector<int> indexMap(n);
    vector<pair<int, int>> sortedNumAndIndexes;

    for (int i = 0; i < n; ++i)
      sortedNumAndIndexes.emplace_back(nums[i], i);

    ranges::sort(sortedNumAndIndexes);

    for (int i = 0; i < n; ++i) {
      const auto& [num, sortedIndex] = sortedNumAndIndexes[i];
      sortedNums.push_back(num);
      indexMap[sortedIndex] = i;
    }

    const int maxLevel = std::bit_width(static_cast<unsigned>(n)) + 1;
    // jump[i][j] := the index of the j-th ancestor of i
    vector<vector<int>> jump(n, vector<int>(maxLevel));

    int right = 0;
    for (int i = 0; i < n; ++i) {
      while (right + 1 < n && sortedNums[right + 1] - sortedNums[i] <= maxDiff)
        ++right;
      jump[i][0] = right;
    }

    for (int level = 1; level < maxLevel; ++level)
      for (int i = 0; i < n; ++i) {
        const int prevJump = jump[i][level - 1];
        jump[i][level] = jump[prevJump][level - 1];
      }

    for (const vector<int>& query : queries) {
      const int u = query[0];
      const int v = query[1];
      const int uIndex = indexMap[u];
      const int vIndex = indexMap[v];
      const int start = min(uIndex, vIndex);
      const int end = max(uIndex, vIndex);
      const int res = minJumps(jump, start, end, maxLevel - 1);
      ans.push_back(res == INT_MAX ? -1 : res);
    }

    return ans;
  }

 private:
  // Returns the minimum number of jumps from `start` to `end` using binary
  // lifting.
  int minJumps(const vector<vector<int>>& jump, int start, int end, int level) {
    if (start == end)
      return 0;
    if (jump[start][0] >= end)
      return 1;
    if (jump[start][level] < end)
      return INT_MAX;
    int j = level;
    for (; j >= 0; --j)
      if (jump[start][j] < end)
        break;
    return (1 << j) + minJumps(jump, jump[start][j], end, j);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class Solution {
  public int[] pathExistenceQueries(int n, int[] nums, int maxDiff, int[][] queries) {
    int[] ans = new int[queries.length];
    int[] indexMap = new int[n];
    int[] sortedNums = new int[n];
    Pair<Integer, Integer>[] sortedNumAndIndexes = new Pair[n];

    for (int i = 0; i < n; ++i)
      sortedNumAndIndexes[i] = new Pair<>(nums[i], i);

    Arrays.sort(sortedNumAndIndexes, Comparator.comparingInt(Pair::getKey));

    for (int i = 0; i < n; ++i) {
      final int num = sortedNumAndIndexes[i].getKey();
      final int sortedIndex = sortedNumAndIndexes[i].getValue();
      sortedNums[i] = num;
      indexMap[sortedIndex] = i;
    }

    final int maxLevel = Integer.SIZE - Integer.numberOfLeadingZeros(n) + 1;
    // jump[i][j] := the index of the j-th ancestor of i
    int[][] jump = new int[n][maxLevel];

    int right = 0;
    for (int i = 0; i < n; ++i) {
      while (right + 1 < n && sortedNums[right + 1] - sortedNums[i] <= maxDiff)
        ++right;
      jump[i][0] = right;
    }

    for (int level = 1; level < maxLevel; ++level)
      for (int i = 0; i < n; ++i) {
        final int prevJump = jump[i][level - 1];
        jump[i][level] = jump[prevJump][level - 1];
      }

    for (int i = 0; i < queries.length; ++i) {
      final int u = queries[i][0];
      final int v = queries[i][1];
      final int uIndex = indexMap[u];
      final int vIndex = indexMap[v];
      final int start = Math.min(uIndex, vIndex);
      final int end = Math.max(uIndex, vIndex);
      final int res = minJumps(jump, start, end, maxLevel - 1);
      ans[i] = res == Integer.MAX_VALUE ? -1 : res;
    }

    return ans;
  }

  // Returns the minimum number of jumps from `start` to `end` using binary
  // lifting.
  private int minJumps(int[][] jump, int start, int end, int level) {
    if (start == end)
      return 0;
    if (jump[start][0] >= end)
      return 1;
    if (jump[start][level] < end)
      return Integer.MAX_VALUE;
    int j = level;
    for (; j >= 0; --j)
      if (jump[start][j] < end)
        break;
    return (1 << j) + minJumps(jump, jump[start][j], end, j);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class Solution:
  def pathExistenceQueries(
      self,
      n: int,
      nums: list[int],
      maxDiff: int,
      queries: list[list[int]],
  ) -> list[int]:
    sortedNumAndIndexes = sorted((num, i) for i, num in enumerate(nums))
    sortedNums = [num for num, _ in sortedNumAndIndexes]
    indexMap = {originalIndex: sortedIndex for sortedIndex,
                (_, originalIndex) in enumerate(sortedNumAndIndexes)}
    maxLevel = n.bit_length() + 1
    # jump[i][j] is the index of the j-th ancestor of i
    jump = [[0] * maxLevel for _ in range(n)]

    right = 0
    for i in range(n):
      while right + 1 < n and sortedNums[right + 1] - sortedNums[i] <= maxDiff:
        right += 1
      jump[i][0] = right

    for level in range(1, maxLevel):
      for i in range(n):
        prevJump = jump[i][level - 1]
        jump[i][level] = jump[prevJump][level - 1]

    def minJumps(start: int, end: int, level: int) -> int:
      """
      Returns the minimum number of jumps from `start` to `end` using binary
      lifting.
      """
      if start == end:
        return 0
      if jump[start][0] >= end:
        return 1
      if jump[start][level] < end:
        return math.inf
      for j in range(level, -1, -1):
        if jump[start][j] < end:
          break
      return (1 << j) + minJumps(jump[start][j], end, j)

    def minDist(u: int, v: int) -> int:
      uIndex = indexMap[u]
      vIndex = indexMap[v]
      start = min(uIndex, vIndex)
      end = max(uIndex, vIndex)
      res = minJumps(start, end, maxLevel - 1)
      return res if res < math.inf else -1

    return [minDist(u, v) for u, v in queries]