数组中{X，Y}对的计数

问题描述

给定一个整数数组 arr 和一个整数 M，计算有多少个有序 {X, Y} 对使得满足以下条件：

• X^Y 中设置位的数量加上 X&Y 中设置位的数量的两倍之和等于 M

其中 ^ 表示按位异或运算符，& 表示按位与运算符。

解题思路

暴力枚举所有的 {X, Y} 对，分别计算 X^Y 和 X&Y 中设置位的数量，统计满足条件的对数。时间复杂度为 $O(n^2)$，无法通过大规模数据测试。

更加高效的解法是使用位运算的性质来统计每一位上的设置位数。具体来说，假设当前处理到的位为第 $i$ 位，那么对于所有的数，统计第 $i$ 位上设置位的数目，并将其分为两类：

• 既在 X^Y 又在 X&Y 中出现的数量，记作 A
• 只在 X^Y 中出现的数量，记作 B
• 只在 X&Y 中出现的数量，记作 C

那么 M 值对第 $i$ 位的贡献就是 2*A + B - C。根据这个公式，我们可以在 $O(n\log V)$ 的时间复杂度内求出满足条件的对数，其中 $V$ 表示数组中元素的最大值。

具体来说，我们可以使用基数排序（radix sort）以及类似二进制分组的技巧来统计每一位上的设置位数。时间复杂度为 $O(n\log V)$。

代码实现

def count_pairs(arr, m):
    """
    计算有多少个有序 `{X, Y}` 对使得满足以下条件：
    - `X^Y` 中设置位的数量加上 `X&Y` 中设置位的数量的两倍之和等于 `M`
    """
    n = len(arr)
    cnt = [0] * 32  # 各位数目计数
    for x in arr:
        # 统计 x 中设置位数目
        for i in range(32):
            cnt[i] += (x >> i) & 1

    # 统计 X^Y 中每一位的数目
    bitcount = [0] * 32  # bitcount[i] 表示第 i 位上不同的数目
    for i in range(32):
        bitcount[i] = (
            sum(cnt[i + 1:]) * 2  # 大于等于 i 的位的数目之和
            + cnt[i] * (n - cnt[i])  # 第 i 位在 X 中设置，不在 Y 中设置
            + (n - cnt[i]) * cnt[i]  # 第 i 位不在 X 中设置，在 Y 中设置
        )

    # 计算 M 中每一位的贡献
    res = 0
    for i in range(32):
        if m & (1 << i):
            res += bitcount[i]

    return res // 2  # 因为 X 和 Y 是有序的，需要除以 2

参考文献

• LeetCode 1865. Finding Pairs With a Certain Sum