1-Trees and Queries (Codeforces 1304E)
Tags: Others, shortest paths

Gildong was hiking a mountain, walking by millions of trees. Inspired by them, he suddenly came up with an interesting idea for trees in data structures: What if we add another edge in a tree?Then he found that such tree-like graphs are called 1-trees. Since Gildong was bored of solving too many tree problems, he wanted to see if similar techniques in trees can be used in 1-trees as well. Instead of solving it by himself, he's going to test you by providing queries on 1-trees.First, he'll provide you a tree (not 1-tree) with $$$n$$$ vertices, then he will ask you $$$q$$$ queries. Each query contains $$$5$$$ integers: $$$x$$$, $$$y$$$, $$$a$$$, $$$b$$$, and $$$k$$$. This means you're asked to determine if there exists a path from vertex $$$a$$$ to $$$b$$$ that contains exactly $$$k$$$ edges after adding a bidirectional edge between vertices $$$x$$$ and $$$y$$$. A path can contain the same vertices and same edges multiple times. All queries are independent of each other; i.e. the added edge in a query is removed in the next query.

Input: The first line contains an integer $$$n$$$ ($$$3 \le n \le 10^5$$$), the number of vertices of the tree.Next $$$n-1$$$ lines contain two integers $$$u$$$ and $$$v$$$ ($$$1 \le u,v \le n$$$, $$$u \ne v$$$) each, which means there is an edge between vertex $$$u$$$ and $$$v$$$. All edges are bidirectional and distinct.Next line contains an integer $$$q$$$ ($$$1 \le q \le 10^5$$$), the number of queries Gildong wants to ask.Next $$$q$$$ lines contain five integers $$$x$$$, $$$y$$$, $$$a$$$, $$$b$$$, and $$$k$$$ each ($$$1 \le x,y,a,b \le n$$$, $$$x \ne y$$$, $$$1 \le k \le 10^9$$$) – the integers explained in the description. It is guaranteed that the edge between $$$x$$$ and $$$y$$$ does not exist in the original tree.

Output: For each query, print "YES" if there exists a path that contains exactly $$$k$$$ edges from vertex $$$a$$$ to $$$b$$$ after adding an edge between vertices $$$x$$$ and $$$y$$$. Otherwise, print "NO".You can print each letter in any case (upper or lower).

Note: The image below describes the tree (circles and solid lines) and the added edges for each query (dotted lines).  Possible paths for the queries with "YES" answers are:   $$$1$$$-st query: $$$1$$$ – $$$3$$$ – $$$2$$$  $$$2$$$-nd query: $$$1$$$ – $$$2$$$ – $$$3$$$  $$$4$$$-th query: $$$3$$$ – $$$4$$$ – $$$2$$$ – $$$3$$$ – $$$4$$$ – $$$2$$$ – $$$3$$$ – $$$4$$$ – $$$2$$$ – $$$3$$$ 

Sample Inputs:
5
1 2
2 3
3 4
4 5
5
1 3 1 2 2
1 4 1 3 2
1 4 1 3 3
4 2 3 3 9
5 2 3 3 9

Sample Outputs:
YES
YES
NO
YES
NO

Solution (Python):
import sys

class RangeQuery:

    def __init__(self, data, func=min):
        self.func = func
        self._data = _data = [list(data)]
        (i, n) = (1, len(_data[0]))
        while 2 * i <= n:
            prev = _data[-1]
            _data.append([func(prev[j], prev[j + i]) for j in range(n - 2 * i + 1)])
            i <<= 1

    def query(self, begin, end):
        depth = (end - begin).bit_length() - 1
        return self.func(self._data[depth][begin], self._data[depth][end - (1 << depth)])

class LCA:

    def __init__(self, root, graph):
        self.time = [-1] * len(graph)
        self.path = []
        self.depth = [0] * len(graph)
        dfs = [root]
        while dfs:
            node = dfs.pop()
            self.path.append(node)
            if self.time[node] == -1:
                self.time[node] = len(self.path) - 1
                for nei in graph[node]:
                    if self.time[nei] == -1:
                        self.depth[nei] = self.depth[node] + 1
                        dfs.append(node)
                        dfs.append(nei)
        self.rmq = RangeQuery((self.time[node] for node in self.path))

    def lca(self, a, b):
        a = self.time[a]
        b = self.time[b]
        if a > b:
            (a, b) = (b, a)
        return self.path[self.rmq.query(a, b + 1)]

def solve(N, edges, queries):
    graph = [[] for i in range(N + 1)]
    for e in edges:
        graph[e[0]].append(e[1])
        graph[e[1]].append(e[0])
    L = LCA(1, graph)

    def lca(u, v):
        return L.lca(u, v)

    def getDist(u, v):
        return L.depth[u] + L.depth[v] - 2 * L.depth[lca(u, v)]
    ans = []
    for (x, y, a, b, k) in queries:
        d_ab = getDist(a, b)
        d_axyb = getDist(a, x) + 1 + getDist(y, b)
        d_ayxb = getDist(a, y) + 1 + getDist(x, b)
        if any((d <= k and (k - d) % 2 == 0 for d in [d_ab, d_axyb, d_ayxb])):
            ans.append('YES')
        else:
            ans.append('NO')
    return '\n'.join(ans)
if __name__ == '__main__':
    N = int(input())
    edges = []
    for i in range(N - 1):
        edges.append(list(map(int, sys.stdin.readline().split())))
    Q = int(input())
    queries = []
    for i in range(Q):
        queries.append(list(map(int, sys.stdin.readline().split())))
    ans = solve(N, edges, queries)
    print(ans)

Solution (Python):
""""""
import sys
inp = sys.stdin.buffer.readline

def input():
    return inp().strip()
out = []

def iin():
    return int(input())

def lin():
    return list(map(int, input().split()))
INF = 10 ** 5 + 7
n = iin()
adj = [[] for i in range(n + 1)]
for _ in range(n - 1):
    (i, j) = lin()
    adj[i].append(j)
    adj[j].append(i)

def dfs(adj, start=1):
    n = len(adj)
    visited = [False] * n
    first = [-1] * n
    euler = []
    height = [-1] * n
    srt = [start]
    height[start] = 1
    parent = [-1] * n
    while srt:
        v = srt.pop()
        if visited[v]:
            euler.append(v)
            continue
        first[v] = len(euler)
        euler.append(v)
        visited[v] = True
        if parent[v] != -1:
            srt.append(parent[v])
        for u in adj[v]:
            if not visited[u]:
                parent[u] = v
                height[u] = height[v] + 1
                srt.append(u)
    return (first, euler, height)
(first, euler, height) = dfs(adj)
euler = [height[i] for i in euler]

class Sparse_Table:

    def __init__(self, arr):
        from math import log
        self.arr = arr
        self.n = len(arr)
        self.logfn = lambda x: log(x, 2)
        self.table = self.build()

    def build(self):
        n = self.n
        sz = int(self.logfn(n)) + 1
        st = [[INF for i in range(sz)] for j in range(n)]
        for i in range(n):
            st[i][0] = self.arr[i]
        for i in range(1, sz):
            step = pow(2, i - 1)
            for j in range(n - step * 2 + 1):
                st[j][i] = min(st[j + step][i - 1], st[j][i - 1])
        return st

    def query(self, begin, end):
        end += 1
        depth = (end - begin).bit_length() - 1
        return min(self.table[begin][depth], self.table[end - (1 << depth)][depth])
sa = Sparse_Table(euler)

def check(i, j):
    (l, r) = (first[i], first[j])
    if l > r:
        (l, r) = (r, l)
    h = sa.query(l, r)
    return height[i] + height[j] - 2 * h
q = iin()
while q:
    q -= 1
    (x, y, a, b, k) = lin()
    ans1 = [check(a, b), check(a, x) + 1 + check(y, b), check(a, y) + 1 + check(x, b)]
    for i in ans1:
        if k - i >= 0 and (k - i) % 2 == 0:
            out.append('YES')
            break
    else:
        out.append('NO')
print('\n'.join(out))

Solution (Python):
import sys, os

class RangeQuery:

    def __init__(self, data, func=min):
        self.func = func
        self._data = _data = [list(data)]
        (i, n) = (1, len(_data[0]))
        while 2 * i <= n:
            prev = _data[-1]
            _data.append([func(prev[j], prev[j + i]) for j in range(n - 2 * i + 1)])
            i <<= 1

    def query(self, begin, end):
        depth = (end - begin).bit_length() - 1
        return self.func(self._data[depth][begin], self._data[depth][end - (1 << depth)])

class LCA:

    def __init__(self, root, graph):
        self.time = [-1] * len(graph)
        self.path = [-1] * len(graph)
        P = [-1] * len(graph)
        t = -1
        dfs = [root]
        while dfs:
            node = dfs.pop()
            self.path[t] = P[node]
            self.time[node] = t = t + 1
            for nei in graph[node]:
                if self.time[nei] == -1:
                    P[nei] = node
                    dfs.append(nei)
        self.rmq = RangeQuery((self.time[node] for node in self.path))

    def __call__(self, a, b):
        if a == b:
            return a
        a = self.time[a]
        b = self.time[b]
        if a > b:
            (a, b) = (b, a)
        return self.path[self.rmq.query(a, b)]
inp = [int(x) for x in sys.stdin.buffer.read().split()]
ii = 0
n = inp[ii]
ii += 1
coupl = [[] for _ in range(n)]
for _ in range(n - 1):
    u = inp[ii] - 1
    ii += 1
    v = inp[ii] - 1
    ii += 1
    coupl[u].append(v)
    coupl[v].append(u)
root = 0
lca = LCA(root, coupl)
depth = [-1] * n
depth[root] = 0
bfs = [root]
for node in bfs:
    for nei in coupl[node]:
        if depth[nei] == -1:
            depth[nei] = depth[node] + 1
            bfs.append(nei)

def dist(a, b):
    c = lca(a, b)
    return depth[a] + depth[b] - 2 * depth[c]
q = inp[ii]
ii += 1
out = []
for _ in range(q):
    x = inp[ii] - 1
    ii += 1
    y = inp[ii] - 1
    ii += 1
    a = inp[ii] - 1
    ii += 1
    b = inp[ii] - 1
    ii += 1
    k = inp[ii]
    ii += 1
    shortest_odd = 10 ** 9 + 11
    shortest_even = 10 ** 9 + 10
    for d in [dist(a, b), dist(a, x) + dist(y, b) + 1, dist(a, y) + dist(x, b) + 1, dist(a, b) + dist(b, x) + dist(y, b) + 1, dist(b, a) + dist(a, x) + dist(y, a) + 1]:
        if d & 1:
            shortest_odd = min(shortest_odd, d)
        else:
            shortest_even = min(shortest_even, d)
    if k & 1:
        out.append(b'YES' if shortest_odd <= k else b'NO')
    else:
        out.append(b'YES' if shortest_even <= k else b'NO')
os.write(1, b'\n'.join(out))