Leetcode 3534. Path Existence Queries in a Graph II

• Leetcode 3534. Path Existence Queries in a Graph II
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3534. Path Existence Queries in a Graph II

1. 解题思路

这一题是题目3532. Path Existence Queries in a Graph I的进阶版本。

在题目3532当中，我们只需要判断query当中的两个元素是否属于同一个簇，而这里，我们还需要额外判断其最短通路的长度。

因此，在构造出DSU进行元素相连性的判断之后，我们还需要额外对元素的最短通路进行一下考察。

此时，我们只需要对所有的元素按照其元素值进行有序排列，然后考察从某一个元素$u$开始依照最大步幅maxDiff进行移动，至少需要几次才能达到目标元素$v$。这个我们可以通过二分搜索进行优化。

2. 代码实现

给出python代码实现如下：

class DSU:
    def __init__(self, N):
        self.root = [i for i in range(N)]
        
    def find(self, k):
        if self.root[k] != k:
            self.root[k] = self.find(self.root[k])
        return self.root[k]
    
    def union(self, a, b):
        x = self.find(a)
        y = self.find(b)
        if x != y:
            self.root[y] = x
        return
    
class Solution:
    def pathExistenceQueries(self, n: int, nums: List[int], maxDiff: int, queries: List[List[int]]) -> List[int]:
        nodes = sorted([(x, i) for i, x in enumerate(nums)])
        dsu = DSU(n)
        for i in range(n-1):
            if nodes[i+1][0] - nodes[i][0] <= maxDiff:
                u, v = nodes[i+1][1], nodes[i][1]
                dsu.union(u, v)
                
        @lru_cache(None)
        def _query(i, j):
            if nodes[i][0] >= nodes[j][0]:
                return 0
            w = nodes[i][1]
            nxt= bisect.bisect_right(nodes, (nums[w]+maxDiff, n)) - 1
            return 1 + _query(nxt, j)

        def query(u, v):
            if nums[u] > nums[v]:
                return query(v, u)
            if dsu.find(u) != dsu.find(v):
                return -1
            elif u == v:
                return 0
            elif nums[u] == nums[v]:
                return 1
            i, j = bisect.bisect_left(nodes, (nums[u], u)), bisect.bisect_left(nodes, (nums[v], v))
            return _query(i, j)
        
        return [query(u, v) for u, v in queries]

提交代码评测得到：耗时1768ms，占用内存175.2MB。