Problem link : Battle of Winterfell
Difficulty : Medium
Setter : Harshit Garg
Code :
from collections import defaultdict
from math import *
from collections import deque
def LcA(a, b):
for i in range(ma, -1, -1):
if lc[a][i] != -1 and lc[a][i] != lc[b][i]:
a = lc[a][i]
b = lc[b][i]
return lc[a][0]
def EQUAL(a, b):
x, y = dis[a], dis[b]
z = abs(x-y)
if z == 0:
return a, b
if x > y:
while z:
p = int(log(z, 2))
a = lc[a][p]
z -= (1 << p)
else:
while z:
p = int(log(z, 2))
b = lc[b][p]
z -= (1 << p)
return a, b
n, qu = map(int, input().split())
d = defaultdict(list)
for i in range(n-1):
a, b = map(int, input().split())
d[a].append(b)
d[b].append(a)
dis = [0 for i in range(n+1)]
q = deque()
q.append(1)
v = [0 for i in range(n+1)]
ma = int(log(n, 2))
lc = [[-1 for i in range(ma+1)] for j in range(n+1)]
while q:
a = q.popleft()
v[a] = 1
for i in d[a]:
if not v[i]:
v[i] = 1
lc[i][0] = a
q.append(i)
dis[i] = dis[a]+1
for j in range(1, ma+1):
for i in range(1, n+1):
if lc[i][j-1] != -1:
p = lc[i][j-1]
lc[i][j] = lc[p][j-1]
ans = []
Q = [list(map(int, input().split())) for i in range(qu)]
for i in Q:
a, b = i
x, y = EQUAL(a, b)
if x == y:
vp = dis[a]+dis[b]-(2*dis[y])
if vp % 2:
ans.append('NO')
else:
ans.append('YES')
continue
z = LcA(x, y)
vp = dis[a]+dis[b]-(2*dis[z])
if vp % 2:
ans.append('NO')
else:
ans.append('YES')
print('\n'.join(ans))