Trees and XOR Queries Again (Codeforces 1902F)
Tags: Math, divide and conquer

You are given a tree consisting of $$$n$$$ vertices. There is an integer written on each vertex; the $$$i$$$-th vertex has integer $$$a_i$$$ written on it.You have to process $$$q$$$ queries. The $$$i$$$-th query consists of three integers $$$x_i$$$, $$$y_i$$$ and $$$k_i$$$. For this query, you have to answer if it is possible to choose a set of vertices $$$v_1, v_2, \dots, v_m$$$ (possibly empty) such that:  every vertex $$$v_j$$$ is on the simple path between $$$x_i$$$ and $$$y_i$$$ (endpoints can be used as well);  $$$a_{v_1} \oplus a_{v_2} \oplus \dots \oplus a_{v_m} = k_i$$$, where $$$\oplus$$$ denotes the bitwise XOR operator. 

Input: The first line contains one integer $$$n$$$ ($$$2 \le n \le 2 \cdot 10^5$$$).The second line contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \le a_i \le 2^{20} - 1$$$).Then $$$n-1$$$ lines follow. Each of them contains two integers $$$u$$$ and $$$v$$$ ($$$1 \le u, v \le n$$$; $$$u \ne v$$$) denoting an edge of the tree.The next line contains one integer $$$q$$$ ($$$1 \le q \le 2 \cdot 10^5$$$) — the number of queries.Then $$$q$$$ lines follow. The $$$i$$$-th of them contains three integers $$$x_i$$$, $$$y_i$$$ and $$$k_i$$$ ($$$1 \le x_i, y_i \le n$$$; $$$0 \le k_i \le 2^{20} - 1$$$).

Output: For each query, print YES if it is possible to form a set of vertices meeting the constraints. Otherwise, print NO.You can print each letter in any case.


Sample Inputs:
4
0 1 2 10
2 1
3 2
4 2
8
3 3 0
3 4 1
3 4 7
1 3 1
1 3 2
1 3 10
1 4 10
1 4 11

Sample Outputs:
YES
YES
NO
YES
YES
NO
YES
YES

Solution (Python):
from __future__ import print_function
from math import *
from collections import defaultdict, deque
import os
import random
import sys
from io import BytesIO, IOBase
from types import GeneratorType

def bootstrap(f, stack=[]):

    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

def main():
    pass
BUFSIZE = 8192

def lcm(a, b):
    return a * b // gcd(a, b)

def ceilDiv(a, b):
    if a % b == 0:
        return a // b
    else:
        return a // b + 1

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')

class myDict:

    def __init__(self, func):
        self.RANDOM = random.randint(0, 1 << 32)
        self.default = func
        self.dict = {}

    def items(self):
        return [(i ^ self.RANDOM, j) for (i, j) in self.dict.items()]

    def __getitem__(self, key):
        myKey = self.RANDOM ^ key
        if myKey not in self.dict:
            self.dict[myKey] = self.default()
        return self.dict[myKey]

    def __setitem__(self, key, item):
        myKey = self.RANDOM ^ key
        self.dict[myKey] = item

    def __iter__(self):
        return iter((i ^ self.RANDOM for i in self.dict))

    def __contains__(self, key):
        return key ^ self.RANDOM in self.dict

    def __repr__(self):
        return repr({i ^ self.RANDOM: j for (i, j) in self.dict.items()})

class mySet:

    def __init__(self):
        self.random = random.randint(0, 1 << 32)
        self.set = set()

    def add(self, key):
        self.set.add(key ^ self.random)

    def remove(self, key):
        self.set.remove(key ^ self.random)

    def __contains__(self, key):
        return key ^ self.random in self.set

    def __iter__(self):
        return iter((i ^ self.random for i in self.set))

    def __len__(self):
        return len(self.set)

    def __repr__(self):
        s = set()
        for item in self.set:
            s.add(item ^ self.random)
        return repr(s)
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')
mod = 10 ** 9 + 7

def inp():
    return int(input())

def inlt():
    return list(map(int, input().split()))

def insr():
    s = input()
    return list(s[:len(s)])

def invr():
    return map(int, input().split())

class XorBasis:

    def __init__(self, bit_length=31):
        self.bit_length = bit_length
        self.basis_size = 0
        self.val = [0] * self.bit_length

    def is_redundant(self, x):
        for i in range(self.bit_length):
            if x >> i & 1:
                if self.val[i] == 0:
                    return False
                x ^= self.val[i]
        return True

    def insert(self, x):
        for i in range(self.bit_length):
            if x >> i & 1:
                if self.val[i] == 0:
                    self.val[i] = x
                    self.basis_size += 1
                    return True
                x ^= self.val[i]
        return False
n = inp()
A = inlt()
l = [[] for i in range(n)]
for i in range(n - 1):
    (a, b) = invr()
    l[a - 1].append(b - 1)
    l[b - 1].append(a - 1)
dp = [[0] * 20 for i in range(n)]
vals = [[] for i in range(n)]
depth = [0 for i in range(n)]
q = deque()
q.append(0)
v = [0] * n
v[0] = 1
if A[0] != 0:
    vals[0].append((A[0], 0))
while q:
    node = q.popleft()
    for child in l[node]:
        if not v[child]:
            v[child] = 1
            q.append(child)
            depth[child] = 1 + depth[node]
            bas = XorBasis(bit_length=20)
            dp[child][0] = node
            if bas.insert(A[child]):
                vals[child].append((A[child], depth[child]))
            for (val, dep) in vals[node]:
                if bas.insert(val):
                    vals[child].append((val, dep))
for j in range(1, 20):
    for i in range(n):
        dp[i][j] = dp[dp[i][j - 1]][j - 1]
for _ in range(inp()):
    (a, b, v) = invr()
    a -= 1
    b -= 1
    bas = XorBasis(bit_length=20)
    rema = a
    remb = b
    if depth[a] > depth[b]:
        (a, b) = (b, a)
    dif = depth[b] - depth[a]
    i = 0
    while dif > 0:
        if dif >> i & 1:
            b = dp[b][i]
            dif ^= 1 << i
        i += 1
    if a == b:
        lca = a
    else:
        for i in range(19, -1, -1):
            if dp[a][i] != dp[b][i]:
                a = dp[a][i]
                b = dp[b][i]
        lca = dp[a][0]
    for (val, dep) in vals[rema]:
        if dep >= depth[lca]:
            bas.insert(val)
    for (val, dep) in vals[remb]:
        if dep >= depth[lca]:
            bas.insert(val)
    if bas.is_redundant(v):
        print('YES')
    else:
        print('NO')

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [-1] * (n + 1)
    for i in q:
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)

def f(u, v):
    return 20 * u + v
n = int(input())
a = [0] + list(map(int, input().split()))
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
(p, dist, parent) = bfs(1)
pow2 = [1]
for _ in range(20):
    pow2.append(2 * pow2[-1])
dp = [-1] * (20 * (n + 1))
for i in range(1, n + 1):
    dp[f(i, 0)] = parent[i]
for i in p:
    k = parent[i]
    for j in range(1, 20):
        if k == -1:
            break
        k = dp[f(k, j - 1)]
        dp[f(i, j)] = k
q = int(input())
(x0, y0, k0) = ([0] * q, [0] * q, [0] * q)
(u0, v0) = ([0] * q, [0] * q)
for q0 in range(q):
    (x, y, k) = map(int, input().split())
    k0[q0] = k
    (u, v) = (x, y)
    if dist[x] > dist[y]:
        (u, v) = (v, u)
    d = dist[v] - dist[u]
    for i in range(20):
        if not d:
            break
        if d & pow2[i]:
            d ^= pow2[i]
            v = dp[f(v, i)]
    for i in range(19, -1, -1):
        if dp[f(u, i)] ^ dp[f(v, i)]:
            (u, v) = (dp[f(u, i)], dp[f(v, i)])
    d = dist[u]
    if u ^ v:
        d -= 1
    (x0[q0], y0[q0]) = (dist[x] - d + 1, dist[y] - d + 1)
    (u0[q0], v0[q0]) = (x, y)
dp0 = dp
dp = [0] * (20 * (n + 1))
for i in range(1, n + 1):
    dp[f(i, 0)] = a[i]
dp1 = [0] * (20 * q)
for i in range(19):
    pi = pow2[i]
    for j in range(q):
        if not pi <= x0[j] < 2 * pi:
            continue
        u = u0[j]
        d = x0[j] - pi
        v = u
        for k in range(20):
            if not d:
                break
            if d & pow2[k]:
                d ^= pow2[k]
                v = dp0[f(v, k)]
        z = []
        for k in range(20 * j, 20 * (j + 1)):
            if not dp1[k]:
                break
            z.append(dp1[k])
        for k in range(20 * u, 20 * (u + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(20 * v, 20 * (v + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(len(z)):
            dp1[f(j, k)] = z[k]
    for j in range(q):
        if not pi <= y0[j] < 2 * pi:
            continue
        u = v0[j]
        d = y0[j] - pi
        v = u
        for k in range(20):
            if not d:
                break
            if d & pow2[k]:
                d ^= pow2[k]
                v = dp0[f(v, k)]
        z = []
        for k in range(20 * j, 20 * (j + 1)):
            if not dp1[k]:
                break
            z.append(dp1[k])
        for k in range(20 * u, 20 * (u + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(20 * v, 20 * (v + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(len(z)):
            dp1[20 * j + k] = z[k]
    for j in p[::-1]:
        k = dp0[f(j, i)]
        if k == -1:
            break
        r = 20 * (j + 1)
        for l in range(20 * j, 20 * (j + 1)):
            if not dp[l]:
                r = l
                break
        for l in range(20 * k, 20 * (k + 1)):
            if r == 20 * (j + 1):
                break
            z = dp[l]
            for w in range(20 * j, r):
                z = min(z, z ^ dp[w])
                if not z:
                    break
            if z:
                dp[r] = z
                r += 1
ans = ['?'] * q
for i in range(q):
    k = k0[i]
    for j in range(20 * i, 20 * (i + 1)):
        k = min(k, k ^ dp1[j])
        if not k:
            break
    ans0 = 'YES' if not k else 'NO'
    ans[i] = ans0
sys.stdout.write('\n'.join(ans))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [-1] * (n + 1)
    for i in q:
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)

def f(u, v):
    return 20 * u + v
n = int(input())
a = [0] + list(map(int, input().split()))
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
(p, dist, parent) = bfs(1)
pow2 = [1]
for _ in range(20):
    pow2.append(2 * pow2[-1])
dp = [-1] * (20 * (n + 1))
for i in range(1, n + 1):
    dp[f(i, 0)] = parent[i]
for i in p:
    k = parent[i]
    for j in range(1, 20):
        if k == -1:
            break
        k = dp[f(k, j - 1)]
        dp[f(i, j)] = k
q = int(input())
(x0, y0, k0) = ([0] * q, [0] * q, [0] * q)
(u0, v0) = ([0] * q, [0] * q)
for q0 in range(q):
    (x, y, k) = map(int, input().split())
    k0[q0] = k
    (u, v) = (x, y)
    if dist[x] > dist[y]:
        (u, v) = (v, u)
    d = dist[v] - dist[u]
    for i in range(20):
        if not d:
            break
        if d & pow2[i]:
            d ^= pow2[i]
            v = dp[f(v, i)]
    for i in range(19, -1, -1):
        if dp[f(u, i)] ^ dp[f(v, i)]:
            (u, v) = (dp[f(u, i)], dp[f(v, i)])
    d = dist[u]
    if u ^ v:
        d -= 1
    (x0[q0], y0[q0]) = (dist[x] - d + 1, dist[y] - d + 1)
    (u0[q0], v0[q0]) = (x, y)
dp0 = dp
dp = [0] * (20 * (n + 1))
for i in range(1, n + 1):
    dp[f(i, 0)] = a[i]
dp1 = [0] * (20 * q)
for i in range(19):
    pi = pow2[i]
    for j in range(q):
        if not pi <= x0[j] < 2 * pi:
            continue
        u = u0[j]
        d = x0[j] - pi
        v = u
        for k in range(20):
            if not d:
                break
            if d & pow2[k]:
                d ^= pow2[k]
                v = dp0[f(v, k)]
        z = []
        for k in range(20 * j, 20 * (j + 1)):
            if not dp1[k]:
                break
            z.append(dp1[k])
        for k in range(20 * u, 20 * (u + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(20 * v, 20 * (v + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(len(z)):
            dp1[f(j, k)] = z[k]
    for j in range(q):
        if not pi <= y0[j] < 2 * pi:
            continue
        u = v0[j]
        d = y0[j] - pi
        v = u
        for k in range(20):
            if not d:
                break
            if d & pow2[k]:
                d ^= pow2[k]
                v = dp0[f(v, k)]
        z = []
        for k in range(20 * j, 20 * (j + 1)):
            if not dp1[k]:
                break
            z.append(dp1[k])
        for k in range(20 * u, 20 * (u + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(20 * v, 20 * (v + 1)):
            if not dp[k]:
                break
            w = dp[k]
            for l in z:
                w = min(w, w ^ l)
                if not w:
                    break
            if w:
                z.append(w)
        for k in range(len(z)):
            dp1[20 * j + k] = z[k]
    for j in p[::-1]:
        k = dp0[f(j, i)]
        if k == -1:
            break
        r = 20 * (j + 1)
        for l in range(20 * j, 20 * (j + 1)):
            if not dp[l]:
                r = l
                break
        for l in range(20 * k, 20 * (k + 1)):
            if r == 20 * (j + 1):
                break
            z = dp[l]
            if not z:
                break
            for w in range(20 * j, r):
                z = min(z, z ^ dp[w])
                if not z:
                    break
            if z:
                dp[r] = z
                r += 1
ans = ['?'] * q
for i in range(q):
    k = k0[i]
    for j in range(20 * i, 20 * (i + 1)):
        k = min(k, k ^ dp1[j])
        if not k:
            break
    ans0 = 'YES' if not k else 'NO'
    ans[i] = ans0
sys.stdout.write('\n'.join(ans))