thảo luận Leetcode + Codeforces, Competitive programming contest. Đường tới Guardian + Candidate Master.

class Solution:
    def findMedian(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        graph = defaultdict(list)
        for f, t, w in edges:
            graph[f].append((t, w))
            graph[t].append((f, w))

        logVal = n.bit_length()

        up = [[-1] * logVal for _ in range(n)]
        depth = [0] * n
        dist = [0] * n

        def dfs(node, parent):
            up[node][0] = parent
            for j in range(1, logVal):
                if up[node][j - 1] != -1:
                    up[node][j] = up[up[node][j - 1]][j - 1]
                else:
                    break
            for child, w in graph[node]:
                if child == parent:
                    continue
                dist[child] = dist[node] + w
                depth[child] = depth[node] + 1
                dfs(child, node)

        dfs(0, -1)

        def lca(u, v):
            if depth[u] < depth[v]:
                u, v = v, u
          
            diff = depth[u] - depth[v]
            for i in range(logVal - 1, -1, -1):
                if diff & (1 << i):
                    u = up[u][i]
          
            if u == v:
                return u
          
            for i in range(logVal - 1, -1, -1):
                if up[u][i] != -1 and up[u][i] != up[v][i]:
                    u = up[u][i]
                    v = up[v][i]
            return up[u][0]

        def lift(u, k):
            for i in range(logVal - 1, -1, -1):
                if u == -1:
                    break
                if k & (1 << i):
                    u = up[u][i]
            return u

        ans = []
        for u, v in queries:
            if u == v:
                ans.append(u)
                continue

            w = lca(u, v)
            totalPathWeight = dist[u] + dist[v] - 2 * dist[w]
            halfTotalWeight = totalPathWeight / 2.0

            distUToW = dist[u] - dist[w]

            # The weighted median node will always be on the path from u to v.
            # We can determine if it's on the u -> LCA segment or LCA -> v segment.

            if halfTotalWeight <= distUToW:
                # The median node is on the path from u to w (LCA)
                # We need to find the node 'x' such that dist(u, x) >= halfTotalWeight.
                # We want the 'x' that is closest to 'u' while satisfying this,
                # meaning we're looking for the smallest 'k' (steps to lift u).
              
                low, high = 0, depth[u] - depth[w]
                stepsToLiftU = high # Initialize with a value that will be updated

                while low <= high:
                    midK = (low + high) // 2
                    liftedNode = lift(u, midK)
                    currentDistFromU = dist[u] - dist[liftedNode]

                    if currentDistFromU >= halfTotalWeight:
                        stepsToLiftU = midK
                        high = midK - 1     
                    else:
                        low = midK + 1       
              
                medianNode = lift(u, stepsToLiftU)

            else:
                # The median node is on the path from w (LCA) to v
                # The cumulative weight from u to w is already 'distUToW'.
                # We need 'halfTotalWeight - distUToW' more distance from w towards v.
                # Let this be 'remainingWeightFromW'.
                # We are looking for node 'x' on w -> v such that dist(w, x) >= remainingWeightFromW.
                # This 'x' is found by lifting 'v' up.
                # We need to find 'k' steps to lift 'v' such that:
                # currentDistFromUToNode = dist(u, w) + dist(w, lift(v, k)) >= halfTotalWeight
                # This simplifies to: distUToW + (dist[lift(v, k)] - dist[w]) >= halfTotalWeight
                # Or, dist[u] + dist[lift(v, k)] - 2 * dist[w] >= halfTotalWeight

                low, high = 0, depth[v] - depth[w]
                stepsToLiftV = 0 # Initialize with 0 steps (node v itself)

                while low <= high:
                    midK = (low + high) // 2
                    liftedNode = lift(v, midK)
                    currentDistFromUToNode = dist[u] + dist[liftedNode] - 2 * dist[w]
                  
                    if currentDistFromUToNode >= halfTotalWeight:
                        stepsToLiftV = midK
                        high = midK - 1       
                    else:
                        low = midK + 1     
              
                # After the loop, `stepsToLiftV` will hold the minimal steps to lift `v`
                # such that `lift(v, stepsToLiftV)` is a candidate for the median.
                # We need the node corresponding to this `stepsToLiftV`.
                medianNode = lift(v, stepsToLiftV)

            ans.append(medianNode)

        return ans