Use this to learn the idea, then write your own version.
1234 5from functools import partial6 789class TreeInfos(object): 10 def __init__(self, children):11 def preprocess(curr, parent, weight):12 if parent != -1:13 W[curr].append(weight)14 P[curr].append(parent) 15 i = 016 while i < len(P[curr]) and i < len(P[P[curr][i]]):17 W[curr].append(max(W[curr][i], W[P[curr][i]][i]))18 P[curr].append(P[P[curr][i]][i])19 i += 120 C[0] += 121 L[curr] = C[0] 22 23 def divide(curr, parent, weight):24 stk.append(partial(postprocess, curr))25 for child, w in reversed(children[curr]):26 if child == parent:27 continue28 stk.append(partial(divide, child, curr, w))29 stk.append(partial(preprocess, curr, parent, weight))30 31 def postprocess(curr):32 R[curr] = C[0] 33 34 N = len(children)35 L, R, P, W, C = [0]*N, [0]*N, [[] for _ in xrange(N)], [[] for _ in xrange(N)], [-1]36 for i in xrange(N):37 if L[i]:38 continue39 stk = []40 stk.append(partial(divide, i, -1, 0))41 while stk:42 stk.pop()()43 self.L, self.R, self.P, self.W = L, R, P, W44 45 def is_ancestor(self, a, b): 46 return self.L[a] <= self.L[b] <= self.R[b] <= self.R[a]47 48 def max_weights(self, a, b):49 def binary_lift(a, b):50 w = 051 for i in reversed(xrange(len(self.P[a]))): 52 if i < len(self.P[a]) and not self.is_ancestor(self.P[a][i], b):53 w = max(w, self.W[a][i])54 a = self.P[a][i]55 return max(w, self.W[a][0])56 57 w = 058 if not self.is_ancestor(a, b):59 w = max(w, binary_lift(a, b))60 if not self.is_ancestor(b, a):61 w = max(w, binary_lift(b, a))62 return w63 64 65class UnionFind(object): 66 def __init__(self, n):67 self.set = range(n)68 self.rank = [0]*n69 70 def find_set(self, x):71 stk = []72 while self.set[x] != x: 73 stk.append(x)74 x = self.set[x]75 while stk:76 self.set[stk.pop()] = x77 return x78 79 def union_set(self, x, y):80 x_root, y_root = map(self.find_set, (x, y))81 if x_root == y_root:82 return False83 if self.rank[x_root] < self.rank[y_root]: 84 self.set[x_root] = y_root85 elif self.rank[x_root] > self.rank[y_root]:86 self.set[y_root] = x_root87 else:88 self.set[y_root] = x_root89 self.rank[x_root] += 190 return True91 92 93class DistanceLimitedPathsExist(object):94 95 def __init__(self, n, edgeList):96 """97 :type n: int98 :type edgeList: List[List[int]]99 """100 edgeList.sort(key = lambda x:x[2])101 self.__uf = UnionFind(n)102 self.__adj = [[] for _ in xrange(n)]103 for index, (i, j, weight) in enumerate(edgeList):104 if not self.__uf.union_set(i, j):105 continue106 self.__adj[i].append((j, weight))107 self.__adj[j].append((i, weight))108 self.__tree_infos = TreeInfos(self.__adj)109 110 def query(self, p, q, limit):111 """112 :type p: int113 :type q: int114 :type limit: int115 :rtype: bool116 """117 if self.__uf.find_set(p) != self.__uf.find_set(q):118 return False119 return self.__tree_infos.max_weights(p, q) < limit120 121 122123124125import collections126import sortedcontainers127import bisect128 129 130class SnapshotArray(object):131 132 def __init__(self, length):133 """134 :type length: int135 """136 self.__snaps = collections.defaultdict(lambda:sortedcontainers.SortedList([(0, 0)]))137 138 def set(self, index, val, snap_id):139 """140 :type index: int141 :type val: int142 :rtype: None143 """144 i = self.__snaps[index].bisect_left((snap_id, float("-inf")))145 if i != len(self.__snaps[index]) and self.__snaps[index][i][0] == snap_id:146 self.__snaps[index].remove(self.__snaps[index][i])147 self.__snaps[index].add((snap_id, val))148 149 def get(self, index, snap_id):150 """151 :type index: int152 :type snap_id: int153 :rtype: int154 """155 i = self.__snaps[index].bisect_left((snap_id+1, float("-inf"))) - 1156 return self.__snaps[index][i][1] 157 158 159class VersionedUnionFind(object): 160 161 def __init__(self, n):162 self.snap_id = 0163 self.set = SnapshotArray(n)164 for i in xrange(n):165 self.set.set(i, i, self.snap_id)166 self.rank = SnapshotArray(n)167 168 def find_set(self, x, snap_id):169 stk = []170 while self.set.get(x, snap_id) != x: 171 stk.append(x)172 x = self.set.get(x, snap_id)173 while stk:174 self.set.set(stk.pop(), x, snap_id)175 return x176 177 def union_set(self, x, y):178 x_root = self.find_set(x, self.snap_id)179 y_root = self.find_set(y, self.snap_id)180 if x_root == y_root:181 return False182 if self.rank.get(x_root, self.snap_id) < self.rank.get(y_root, self.snap_id): 183 self.set.set(x_root, y_root, self.snap_id)184 elif self.rank.get(x_root, self.snap_id) > self.rank.get(y_root, self.snap_id):185 self.set.set(y_root, x_root, self.snap_id)186 else:187 self.set.set(y_root, x_root, self.snap_id)188 self.rank.set(x_root, self.rank.get(x_root, self.snap_id)+1, self.snap_id)189 return True190 191 def snap(self):192 self.snap_id += 1193 194 195class DistanceLimitedPathsExist2(object):196 197 def __init__(self, n, edgeList):198 """199 :type n: int200 :type edgeList: List[List[int]]201 """202 edgeList.sort(key = lambda x:x[2])203 self.__uf = VersionedUnionFind(n)204 self.__weights = []205 for index, (i, j, weight) in enumerate(edgeList):206 if not self.__uf.union_set(i, j):207 continue208 self.__uf.snap()209 self.__weights.append(weight) 210 211 def query(self, p, q, limit):212 """213 :type p: int214 :type q: int215 :type limit: int216 :rtype: bool217 """218 snap_id = bisect.bisect_left(self.__weights, limit)-1219 if snap_id == -1:220 return False221 return self.__uf.find_set(p, snap_id) == self.__uf.find_set(q, snap_id)222