Editorial for Pondo Tree Distance

Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Approach

Root the tree at vertex $1$ ~1~ and compute the depth of every vertex. The unique path between $u$ ~u~ and $v$ ~v~ goes up from $u$ ~u~ to $\mathrm{lca}(u, v)$ ~\mathrm{lca}(u, v)~ and down to $v$ ~v~, so

$\mathrm{dist}(u, v) = \mathrm{depth}[u] + \mathrm{depth}[v] - 2 \cdot \mathrm{depth}[\mathrm{lca}(u, v)]$ .

Compute lowest common ancestors the same way as in Pondo LCA II: Euler tour of visits, first-occurrence indices, and a sparse table on tour depths. Overall time is $O((n + q) \log n)$ ~O((n + q) \log n)~ to build the table and $O(n + q)$ ~O(n + q)~ to answer the queries.

Solution (Python)

import sys

sys.setrecursionlimit(1_000_000)


class SparseTable:
    def __init__(self, arr):
        self.arr = arr
        self.n = len(arr)
        self.log = [0] * (self.n + 1)
        self._compute_logs()
        self.st = self._build_sparse_table()

    def _compute_logs(self):
        for i in range(2, self.n + 1):
            self.log[i] = self.log[i // 2] + 1

    def _build_sparse_table(self):
        k = self.log[self.n] + 1
        st = [[(0, 0)] * k for _ in range(self.n)]
        for i in range(self.n):
            st[i][0] = (self.arr[i], i)
        j = 1
        while (1 << j) <= self.n:
            i = 0
            while i + (1 << j) - 1 < self.n:
                if st[i][j - 1][0] < st[i + (1 << (j - 1))][j - 1][0]:
                    st[i][j] = st[i][j - 1]
                else:
                    st[i][j] = st[i + (1 << (j - 1))][j - 1]
                i += 1
            j += 1
        return st

    def query(self, L, R):
        j = self.log[R - L + 1]
        left = self.st[L][j]
        right = self.st[R - (1 << j) + 1][j]
        if left[0] < right[0]:
            return left
        return right


data = list(map(int, sys.stdin.buffer.read().split()))
it = iter(data)
n = next(it)
q = next(it)
adj = [[] for _ in range(n)]
for _ in range(n - 1):
    u = next(it) - 1
    v = next(it) - 1
    adj[u].append(v)
    adj[v].append(u)

tour = []
tour_h = []
first = [-1] * n
depth = [0] * n


def dfs(u, p, h):
    depth[u] = h
    first[u] = len(tour)
    tour.append(u)
    tour_h.append(h)
    for v in adj[u]:
        if v == p:
            continue
        dfs(v, u, h + 1)
        tour.append(u)
        tour_h.append(h)


dfs(0, -1, 0)
st = SparseTable(tour_h)

out = []
for _ in range(q):
    u = next(it) - 1
    v = next(it) - 1
    l = first[u]
    r = first[v]
    if l > r:
        l, r = r, l
    _, idx = st.query(l, r)
    lca = tour[idx]
    out.append(str(depth[u] + depth[v] - 2 * depth[lca]))

print("\n".join(out))

Comments

There are no comments at the moment.