计算给定数组中按位XOR小于K的对

一、问题描述

给定一个整数数组 nums 和一个整数 k ，请计算该数组中有多少对元素对 (i, j) 满足 i<j 且 nums[i] XOR nums[j] < k 。

二、示例

示例 1:

输入: nums = [1, 4, 2, 7], k = 5
输出: 6
解释: 一共有 6 对符合题意的元素对：
     (1, 2), (1, 4), (1, 7), (2, 4), (2, 7), (4, 7)

示例 2:

输入: nums = [9, 8, 7, 6, 5], k = 5
输出: 2
解释: 一共有 2 对符合题意的元素对：
     (5, 6), (5, 7)

三、思路

暴力求解: 遍历数组中的每一个元素，再遍历其后面的元素，计算两者的 XOR 值，如果小于给定的 k ，则计数器加一。

时间复杂度: $O(n^2)$

优化: 利用 Trie 树。

对于每个数字，我们从高到低依次计算二进制位上是否存在某些位可以和后面的数字与起来小于 k 。我们从最高位往下计算，如果两个数字在某一位上的值都为 0 ，那么它们与后面数字的 XOR 值在这一位上也一定为 0 ，此时我们需要寻找后面的数字在这一位上是否存在值为 1 的二进制位，这样才有可能使得它们的 XOR 值小于 k 。如果后面数字在这一位上也为 0 ，那么接下来我们只需要考虑更低的二进制位就好了。如果后面数字在这一位上为 1 ，那么它们的 XOR 值在这一位上一定是 1 ，而前面数字在这一位上为 0 ，可以表示为节点的左儿子与右儿子在这一位上的取值不同，因此我们可以在 Trie 树上往右走，将当前二进制位的 XOR 值加入到答案中，然后考虑更低的二进制位，继续按照上述方法遍历 Trie 树，最终得到答案。

时间复杂度: $O(32n)$

四、图示解释

Trie 树的结构如下所示：

我们假设输入的 nums 数组为 [1, 4, 2, 7] ， k 为 5 。我们按照上述方法在 Trie 树上遍历，遍历过程如下所示：

最终得到的答案为 6。

五、代码实现

暴力求解:

def countPairs(nums: List[int], k: int) -> int:
    n = len(nums)
    res = 0
    for i in range(n):
        for j in range(i + 1, n):
            if nums[i] ^ nums[j] < k:
                res += 1
    return res

利用 Trie 树:

class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.count = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, num: int):
        node = self.root
        for i in range(31, -1, -1):
            bit = (num >> i) & 1
            node = node.children[bit]
            node.count += 1

    def query(self, num: int, k: int) -> int:
        res, node = 0, self.root
        for i in range(31, -1, -1):
            if not node:
                return res
            bit1, bit2 = (num >> i) & 1, (k >> i) & 1
            if bit2 == 1:
                if bit1 == 0:
                    res += node.children[1].count
                node = node.children.get(0)
            else:
                if bit1 == 1:
                    res += node.children[0].count
                node = node.children.get(1)
        return res

def countPairs(nums: List[int], k: int) -> int:
    trie = Trie()
    res = 0
    for num in nums:
        trie.insert(num)
        res += trie.query(num, k)
    return res

返回的代码片段：

暴力求解:

```python
def countPairs(nums: List[int], k: int) -> int:
    n = len(nums)
    res = 0
    for i in range(n):
        for j in range(i + 1, n):
            if nums[i] ^ nums[j] < k:
                res += 1
    return res

利用 Trie 树:

class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.count = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, num: int):
        node = self.root
        for i in range(31, -1, -1):
            bit = (num >> i) & 1
            node = node.children[bit]
            node.count += 1

    def query(self, num: int, k: int) -> int:
        res, node = 0, self.root
        for i in range(31, -1, -1):
            if not node:
                return res
            bit1, bit2 = (num >> i) & 1, (k >> i) & 1
            if bit2 == 1:
                if bit1 == 0:
                    res += node.children[1].count
                node = node.children.get(0)
            else:
                if bit1 == 1:
                    res += node.children[0].count
                node = node.children.get(1)
        return res

def countPairs(nums: List[int], k: int) -> int:
    trie = Trie()
    res = 0
    for num in nums:
        trie.insert(num)
        res += trie.query(num, k)
    return res