Bakry and Partitioning (Codeforces 1592C)
Tags: Others, trees

Bakry faced a problem, but since he's lazy to solve it, he asks for your help.You are given a tree of $$$n$$$ nodes, the $$$i$$$-th node has value $$$a_i$$$ assigned to it for each $$$i$$$ from $$$1$$$ to $$$n$$$. As a reminder, a tree on $$$n$$$ nodes is a connected graph with $$$n-1$$$ edges.You want to delete at least $$$1$$$, but at most $$$k-1$$$ edges from the tree, so that the following condition would hold:For every connected component calculate the bitwise XOR of the values of the nodes in it. Then, these values have to be the same for all connected components.Is it possible to achieve this condition?

Input: Each test contains multiple test cases. The first line contains the number of test cases $$$t$$$ $$$(1 \leq t \leq 5 \cdot 10^4)$$$. Description of the test cases follows.The first line of each test case contains two integers $$$n$$$ and $$$k$$$ $$$(2 \leq k \leq n \leq 10^5)$$$.The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, ..., a_n$$$ $$$(1 \leq a_i \leq 10^9)$$$.The $$$i$$$-th of the next $$$n-1$$$ lines contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$, $$$u_i\neq v_i$$$), which means that there's an edge between nodes $$$u_i$$$ and $$$v_i$$$.It is guaranteed that the given graph is a tree.It is guaranteed that the sum of $$$n$$$ over all test cases doesn't exceed $$$2 \cdot 10^5$$$.

Output: For each test case, you should output a single string. If you can delete the edges according to the conditions written above, output "YES" (without quotes). Otherwise, output "NO" (without quotes).You can print each letter of "YES" and "NO" in any case (upper or lower).

Note: It can be shown that the objection is not achievable for first, third, and fifth test cases.In the second test case, you can just remove all the edges. There will be $$$5$$$ connected components, each containing only one node with value $$$3$$$, so the bitwise XORs will be $$$3$$$ for all of them.In the fourth test case, this is the tree: .You can remove an edge $$$(4,5)$$$The bitwise XOR of the first component will be, $$$a_1 \oplus a_2 \oplus a_3 \oplus a_4 = 1 \oplus 6 \oplus 4 \oplus 1 = 2$$$ (where $$$\oplus$$$ denotes the bitwise XOR). The bitwise XOR of the second component will be, $$$a_5 = 2$$$. 

Sample Inputs:
5
2 2
1 3
1 2
5 5
3 3 3 3 3
1 2
2 3
1 4
4 5
5 2
1 7 2 3 5
1 2
2 3
1 4
4 5
5 3
1 6 4 1 2
1 2
2 3
1 4
4 5
3 3
1 7 4
1 2
2 3

Sample Outputs:
NO
YES
NO
YES
NO

Solution (Python):
def main():
    alpha = 'abcdefghijklmnopqrstuvwxyz'
    ALPHA = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    inf = 1e+17
    mod = 10 ** 9 + 7

    def factorial(n):
        f = 1
        for i in range(1, n + 1):
            f = f * i % mod
        return f

    def factorial_by(n, by):
        f = 1
        for i in range(1, n + 1):
            if i == by:
                continue
            f = f * i % mod
        return f

    def ncr(n, r):
        num = den = 1
        for i in range(r):
            num = num * (n - i) % mod
            den = den * (i + 1) % mod
        return num * pow(den, mod - 2, mod) % mod

    def primeFactors(num):
        pf = []
        while num % 2 == 0:
            pf.append(2)
            num = num // 2
        for i in range(3, int(math.sqrt(num)) + 1, 2):
            while num % i == 0:
                pf.append(i)
                num = num // i
        if num > 2:
            pf.append(num)
        return pf

    class Node(object):

        def __init__(self, name, d):
            self.d = d
            self.name = name

        def __repr__(self):
            pass

    def dfs(node, n, tree, a, taken, dp, XOR, counter):
        taken.add(node)
        xor = a[node]
        for child in tree[node]:
            if child not in taken:
                xor ^= dfs(child, n, tree, a, taken, dp, XOR, counter)
        if xor == XOR:
            counter[0] += 1
            xor = 0
        dp[node] = xor
        return dp[node]

    def solve(n, k, tree, a):
        a = [0] + a
        XOR = 0
        for _ in range(1, n + 1):
            XOR ^= a[_]
        if XOR:
            dp = [0] * (n + 1)
            if k > 2:
                counter = [0]
                taken = set()
                dfs(1, n, tree, a, taken, dp, XOR, counter)
                if counter[0] >= 3:
                    return 'YES'
                else:
                    return 'NO'
            else:
                return 'NO'
        return 'YES'
    t = int(input())
    ans = []
    for _ in range(t):
        (n, k) = map(int, input().split())
        a = [int(x) for x in input().split()]
        tree = []
        degrees = [0] * (n + 1)
        for i in range(n + 1):
            tree.append([])
        for i in range(n - 1):
            (u, v) = map(int, input().split())
            tree[u].append(v)
            tree[v].append(u)
        ans.append(solve(n, k, tree, a))
    p = 1
    for answer in ans:
        print(answer)
        p += 1
if __name__ == '__main__':
    import sys, threading
    import bisect
    import math
    import itertools
    from sys import stdout
    import heapq
    from queue import PriorityQueue
    from collections import deque
    sys.setrecursionlimit(2 ** 32 // 2 - 1)
    threading.stack_size(1 << 27)
    input = sys.stdin.readline
    thread = threading.Thread(target=main)
    thread.start()
    thread.join()

Solution (Python):
import sys, threading, collections

def solve():
    from collections import defaultdict
    input = sys.stdin.readline

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

    def inps():
        return input().strip()

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

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

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

    def pre_compute_xor(i):
        global cnt
        visited[i] = 1
        x = l[i]
        for j in d[i]:
            if visited[j] == 0:
                visited[j] = 1
                x ^= pre_compute_xor(j)
        if x == xr:
            cnt += 1
            return 0
        xor_values[i] = x
        return x
    for _ in range(inp()):
        (n, k) = invr()
        l = inlt()
        xr = 0
        global cnt
        cnt = 0
        xor_values = [0] * n
        d = defaultdict(lambda : [])
        for i in range(n - 1):
            (a, b) = invr()
            a -= 1
            b -= 1
            d[a].append(b)
            d[b].append(a)
        for i in range(n):
            xr = xr ^ l[i]
        if xr == 0:
            print('YES')
        elif k < 3:
            print('NO')
        else:
            cnt = 0
            visited = [0] * n
            pre_compute_xor(1)
            if cnt >= 3:
                print('YES')
            else:
                print('NO')
cnt = 0
sys.setrecursionlimit(10 ** 9)
threading.stack_size(10 ** 8)
threading.Thread(target=solve).start()

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

def solve():
    input = sys.stdin.readline

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

    def inps():
        return input().strip()

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

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

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

    def pre_compute_xor(i):
        global cnt
        visited[i] = 1
        x = l[i]
        for j in d[i]:
            if visited[j] == 0:
                visited[j] = 1
                x ^= pre_compute_xor(j)
        if x == xr:
            cnt += 1
            return 0
        xor_values[i] = x
        return x
    for _ in range(inp()):
        (n, k) = invr()
        l = inlt()
        xr = 0
        global cnt
        cnt = 0
        xor_values = [0] * n
        d = defaultdict(lambda : [])
        for i in range(n - 1):
            (a, b) = invr()
            a -= 1
            b -= 1
            d[a].append(b)
            d[b].append(a)
        for i in range(n):
            xr = xr ^ l[i]
        if xr == 0:
            print('YES')
        elif k < 3:
            print('NO')
        else:
            cnt = 0
            visited = [0] * n
            pre_compute_xor(1)
            if cnt >= 3:
                print('YES')
            else:
                print('NO')
cnt = 0
import collections, sys, threading
sys.setrecursionlimit(10 ** 5)
threading.stack_size(10 ** 8)
threading.Thread(target=solve).start()