Problem statement

https://leetcode.com/problems/checking-existence-of-edge-length-limited-paths-ii/

Solution

This is an online version of problem 1697: the different is that there we were given all queries in advance and hence we were able to sort them. Here it is not the case and we need to think about something different.

We can notice that if we consider all paths between two nodes u and v, then the path which minimizes the maximum value of edges in this path always belong to minumim spanning tree! The idea is that when we construct MST, we add edges starting with the smallest weights and do not add edges if they are inside one connected component. So, the plan is the following:

  1. Construct MST, using kruskal algorithm, write it to self.MST object.
  2. What we have now is spanning tree (or forest) and we need to find the weight of minimum edge which connects given nodes u and v. For this we need to use binary lifting technique: we need to keep several pieces of information:
  3. self.depth are depths of nodes in our tree (or forest).
  4. self.parents[node][level] is ancestor of rank 1<<level of node.
  5. self.max_path[node][level] is maximum of all values on path between node and its 1<<level ancestor.
  6. Once we found MST, we traverse it using usual dfs and evaluate depths, and initialize max_path and parents.
  7. After this we fill self.parents and self.max_path, using parent of parent trick. We denote parent of root as -1 for more simpler navigation.
  8. Finally, how query(p, q, limit) works: we check if p and q are in the same connected component, and if not, return False. If they are, then we need to run LCA(u, v) function.
  9. LCU(u, v) will find the minimum weight on path between u and v. First, we need to list v such that it has the same depth as u. Then we need to go up with decreasing powers of 2. Imagine, that we have distance 21 until LCA. Then we start with bin power of 2: 1<<(bits - 1) and ask question: are these ansectors are the same. Imagine bits = 7, then we ask first if 1<<6 ansectors are the same? Yes, decrease to 1<<5 then 1<<4. For these two values ther are not the same, so we make distance = 21 - 16 = 5 now. We will stop just before our LCA, so in the end we need to do one more iteration.

Complexity

It is O(E log E + n) to build MST and then O(n) to run dfs. Also we have O(n log n) to build our parents and max_path. Finally we have Q queries with O(log n) complexity for each. So, total time complexity is O(E log E + n log n + Q log n). Space complexity is O(n + Q).

Code

class DSU:
    def __init__(self, N):
        self.p = list(range(N))

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

    def union(self, x, y):
        xr = self.find(x)
        yr = self.find(y)
        self.p[xr] = yr

class DistanceLimitedPathsExist:
    def __init__(self, n, E):
        self.bits = bits = ceil(log2(n))
        self.E = sorted(E, key = lambda x:x[2])  
        self.visited = [0]*n
        self.MST = defaultdict(list)
        self.depths = [0]*n
        self.dsu = DSU(n)
        self.parents = [[-1]*bits for _ in range(n)]
        self.max_path = [[0]*bits for _ in range(n)]
        self.kruskal()

        for i in range(n):
            if self.visited[i] == 0:
                self.depths[i] = 1
                self.dfs(i)
                self.parents[i][0] = i

        for i in range(1, bits):
            for j in range(n):
                pp = self.parents[j][i-1]
                self.parents[j][i] = self.parents[pp][i-1]
                self.max_path[j][i] = max(self.max_path[j][i-1], self.max_path[pp][i-1])

    def query(self, p, q, limit):
        return False if self.dsu.find(p) != self.dsu.find(q) else self.LCA(p, q) < limit

    def LCA(self, u, v):
        if self.depths[u] > self.depths[v]:
            u, v = v, u
        
        diff = self.depths[v] - self.depths[u] 

        weight = 0
        places = [i for i in range(diff.bit_length()) if diff&(1<<i) != 0]
        for index in places:
            weight = max(weight, self.max_path[v][index])
            v = self.parents[v][index]

        if u == v: return weight
        for i in range(self.bits - 1, -1, -1):
            if self.parents[u][i] != self.parents[v][i]:
                weight = max([weight, self.max_path[u][i], self.max_path[v][i]])
                u = self.parents[u][i]
                v = self.parents[v][i]

        return max([weight, self.max_path[u][0], self.max_path[v][0]])

    def kruskal(self):
        for u, v, w in self.E:
            if self.dsu.find(u) == self.dsu.find(v): continue
            self.dsu.union(u, v)
            self.MST[u].append((v, w))
            self.MST[v].append((u, w))

    def dfs(self, u):
        self.visited[u] = 1
        for v, w in self.MST[u]:
            if self.visited[v] == 0:
                self.depths[v] = self.depths[u] + 1
                self.parents[v][0] = u
                self.max_path[v][0] = w
                self.dfs(v)