Chương trình tính tổng KPR cho tất cả truy vấn trong Python

Mô tả bài toán

Cho mảng nums và danh sách truy vấn queries[i] = [k, p, r]. Tính kpr_sum = sum_{i=p}^{r-1} sum_{j=i+1}^{r} (K XOR (A[i] XOR A[j])). Trả về modulo 10^9+7.

Ví dụ

nums = [1,2,3], queries = [[1,1,3],[2,1,3]] -> [5, 4].

Giải thuật (Bitwise Prefix Sum)

1. Tiền xử lý: với mỗi bit 0..19, tính prefix sum số bit 1.
2. Với mỗi truy vấn: với mỗi bit, tính số cặp có XOR bit = 0/1, kết hợp với bit của K để tính đóng góp.
3. Cộng dồn theo trọng số 2^i.

Triển khai Python

def solve(nums, queries):
    m = 10**9 + 7
    N = len(nums)
    q_cnt = len(queries)
    C = []
    for i in range(20):
        R = [0]
        t = 0
        for x in nums:
            t += (x >> i) & 1
            R.append(t)
        C.append(R)
    res = []
    for j in range(q_cnt):
        K, P, R = queries[j]
        d = R - P + 1
        t = 0
        for i in range(20):
            n1 = C[i][R] - C[i][P - 1]
            n0 = d - n1
            if (K >> i) & 1:
                x = (n1 * (n1 - 1) + n0 * (n0 - 1)) >> 1
            else:
                x = n1 * n0
            t = (t + (x << i)) % m
        res.append(t)
    return res

nums = [1, 2, 3]
queries = [[1, 1, 3], [2, 1, 3]]
print(solve(nums, queries))

Đầu vào

[1,2,3], [[1,1,3],[2,1,3]]

Đầu ra

[5, 4]