Keep XOR Low (Codeforces 1616H)
Tags: Math, divide and conquer

You are given an array $$$a_1, a_2, \ldots, a_n$$$ and an integer $$$x$$$.Find the number of non-empty subsets of indices of this array $$$1 \leq b_1 < b_2 < \ldots < b_k \leq n$$$, such that for all pairs $$$(i, j)$$$ where $$$1 \leq i < j \leq k$$$, the inequality $$$a_{b_i} \oplus a_{b_j} \leq x$$$ is held. Here, $$$\oplus$$$ denotes the bitwise XOR operation. As the answer may be very large, output it modulo $$$998\,244\,353$$$.

Input: The first line of the input contains two integers $$$n$$$ and $$$x$$$ ($$$1 \leq n \leq 150\,000$$$, $$$0 \leq x < 2^{30}$$$). Here, $$$n$$$ is the size of the array.The next line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$0 \leq a_i < 2^{30}$$$): the array itself.

Output: Print one integer: the number of non-empty subsets such that the bitwise XOR of every pair of elements is at most $$$x$$$, modulo $$$998\,244\,353$$$.


Sample Inputs:
4 2
0 1 2 3

3 6
4 2 2

4 0
1 1 2 2

Sample Outputs:
8

7

6

Solution (Python):
limite = 0
v = []

def contaSoloRipetizioni(v):
    i = 0
    s = 0
    while i < len(v):
        molt = 1
        while 1:
            i += 1
            if i == len(v) or v[i] != v[i - 1]:
                break
            molt += 1
        s += (1 << molt) - 1
    return s

def findBorder(v, i1, i2, x):
    if v[i1] < x:
        return i1
    if v[i2 - 1] >= x:
        return i2
    while 1:
        k = (i1 + i2 - 1) // 2
        if k == i1:
            return i2 - 1
        if v[k] < x:
            i2 = k + 1
        else:
            i1 = k

def calcola(v, i1, i2, M, bit):
    k = findBorder(v, i1, i2, M)
    C_C = (1 << k - i1) - 1
    C_D = (1 << i2 - k) - 1
    if i1 == k or i2 == k:
        C_CD = 0
    else:
        C_CD = almenoUnoInEntrambi(v, i1, k, M + bit // 2, k, i2, M - bit // 2, bit // 2)
    return C_C + C_D + C_CD + 1

def almenoUnoInEntrambi(v, c1, c2, Mc, d1, d2, Md, bit):
    cL = findBorder(v, c1, c2, Mc)
    dL = findBorder(v, d1, d2, Md)
    if bit == 0:
        print('IMPOSSIBILE qui!')
        exit()
    if limite & bit == 0:
        if c2 == cL or d2 == dL:
            q1 = 0
        elif v[cL] == v[c2 - 1] and v[dL] == v[d2 - 1]:
            q1 = ((1 << c2 - cL) - 1) * ((1 << d2 - dL) - 1) if v[cL] ^ v[dL] <= limite else 0
        else:
            q1 = almenoUnoInEntrambi(v, cL, c2, Mc - bit // 2, dL, d2, Md - bit // 2, bit // 2)
        if c1 == cL or d1 == dL:
            q2 = 0
        elif v[c1] == v[cL - 1] and v[d1] == v[dL - 1]:
            q2 = ((1 << cL - c1) - 1) * ((1 << dL - d1) - 1) if v[c1] ^ v[d1] <= limite else 0
        else:
            q2 = almenoUnoInEntrambi(v, c1, cL, Mc + bit // 2, d1, dL, Md + bit // 2, bit // 2)
        return q1 + q2
    else:
        if c2 == cL or d1 == dL:
            q1 = 0
        elif v[cL] == v[c2 - 1] and v[d1] == v[dL - 1]:
            q1 = ((1 << c2 - cL) - 1) * ((1 << dL - d1) - 1) if v[cL] ^ v[d1] <= limite else 0
        else:
            q1 = almenoUnoInEntrambi(v, cL, c2, Mc - bit // 2, d1, dL, Md + bit // 2, bit // 2)
        if c1 == cL or d2 == dL:
            q2 = 0
        elif v[c1] == v[cL - 1] and v[dL] == v[d2 - 1]:
            q2 = ((1 << cL - c1) - 1) * ((1 << d2 - dL) - 1) if v[c1] ^ v[dL] <= limite else 0
        else:
            q2 = almenoUnoInEntrambi(v, c1, cL, Mc + bit // 2, dL, d2, Md - bit // 2, bit // 2)
        v1 = ((1 << c2 - cL) - 1) * ((1 << d2 - dL) - 1)
        v2 = q2 * ((1 << dL - d1) - 1 + ((1 << c2 - cL) - 1) + 1)
        v3 = q1 * ((1 << d2 - dL) - 1 + ((1 << cL - c1) - 1) + 1)
        v4 = ((1 << cL - c1) - 1) * ((1 << dL - d1) - 1)
        return v1 + v2 + v3 + v4 + q1 * q2

def solve():
    v.sort(reverse=True)
    if limite == 0:
        return contaSoloRipetizioni(v)
    if len(v) == 1:
        return 1
    M = 1 << limite.bit_length() - 1
    j = 2
    Ctot = 0
    i2 = len(v)
    while 1:
        i = findBorder(v, 0, i2, j * M)
        if i == i2:
            C_B = 0
        else:
            C_B = calcola(v, i, i2, j * M - M, M) - 1
        Ctot = (Ctot + C_B) % 998244353
        if i == 0:
            break
        i2 = i
        j += 2
    return Ctot
n = 1
for i in range(n):
    (x, limite) = map(int, input().split())
    (*v,) = map(int, input().split())
    print(solve())

Solution (Python):
import sys
I = lambda : [*map(int, sys.stdin.readline().split())]
M = 998244353
(n, x) = I()
a = I()
bits = [x >> i & 1 for i in range(30)]

def f(a, b, i):
    if i == -1:
        return (pow(2, len(a), M) - 1) * (pow(2, len(b), M) - 1) % M
    if len(a) == 0 or len(b) == 0:
        return 0
    a0 = []
    a1 = []
    b0 = []
    b1 = []
    for guy in a:
        if guy >> i & 1:
            a1.append(guy)
        else:
            a0.append(guy)
    for guy in b:
        if guy >> i & 1:
            b1.append(guy)
        else:
            b0.append(guy)
    if bits[i]:
        return ((f(a0, b1, i - 1) + pow(2, len(a0), M) + pow(2, len(b1), M) - 1) * (f(a1, b0, i - 1) + pow(2, len(a1), M) + pow(2, len(b0), M) - 1) - pow(2, len(a), M) - pow(2, len(b), M) + 1) % M
        return x
    else:
        return (f(a0, b0, i - 1) + f(a1, b1, i - 1)) % M
        return x

def g(a, i):
    if len(a) == 0:
        return 0
    if i == -1:
        return pow(2, len(a), M) - 1
    else:
        newa = []
        newb = []
        for guy in a:
            if guy >> i & 1:
                newa.append(guy)
            else:
                newb.append(guy)
        if bits[i] == 0:
            return (g(newa, i - 1) + g(newb, i - 1)) % M
        else:
            return (pow(2, len(newa), M) + pow(2, len(newb), M) + f(newa, newb, i - 1) - 2) % M
print(g(a, 29))

Solution (Python):
I = lambda : [*map(int, input().split())]
M = 998244353
(n, X) = I()
a = I()
q = lambda x: pow(2, x, M) - 1
E = len

def f(a, b, i):
    if i < 0 or not a or (not b):
        return q(E(a)) * q(E(b))
    (w, x, y, z) = ([], [], [], [])
    for g in a:
        if g >> i & 1:
            w.append(g)
        else:
            x.append(g)
    for g in b:
        if g >> i & 1:
            y.append(g)
        else:
            z.append(g)
    if X >> i & 1:
        return (f(x, y, i - 1) + q(E(x)) + q(E(y)) + 1) * (f(w, z, i - 1) + q(E(w)) + q(E(z)) + 1) - q(E(a)) - q(E(b)) - 1
    else:
        return f(x, z, i - 1) + f(w, y, i - 1)

def F(a, i):
    if i < 0 or not a:
        return q(E(a))
    else:
        (x, y) = ([], [])
        for g in a:
            if g >> i & 1:
                x.append(g)
            else:
                y.append(g)
        if X >> i & 1:
            return q(E(x)) + q(E(y)) + f(x, y, i - 1)
        else:
            return F(x, i - 1) + F(y, i - 1)
print(F(a, 29) % M)