LeetCode题解(1803)：统计异或值在范围内的数对有多少(Python)

题目：​​原题链接​​（困难）

标签：字典树

解法一：

class Solution:
    """
    字典树
    时间复杂度：O(NlogN)
    """

    class _Node:
        __slots__ = "num", "left", "right"

        def __init__(self):
            self.num = 0  # 当前树包含元素数量
            self.left = None  # 左子树：当前位为0
            self.right = None  # 右子树：当前位为1

        @property
        def right_num(self):
            return self.right.num if self.right else 0

        @property
        def left_num(self):
            return self.left.num if self.left else 0

    def countPairs(self, nums: List[int], low: int, high: int) -> int:
        # 计算最大位数
        max_bit_length = max(num.bit_length() for num in nums)

        # 调整最大值：若最大值大于最大位数，则将最大值调为当前最大位数全为1的情况
        if high.bit_length() > max_bit_length:
            high = 2 ** max_bit_length - 1

        # ---------- 构造关于二进制位的字典树 ----------
        tree = self._Node()

        for num in nums:
            node = tree
            for k in range(max_bit_length - 1, -1, -1):
                if (num >> k) & 1:  # 当前位为1
                    if node.right is None:
                        node.right = self._Node()
                    node = node.right
                else:  # 当前位为0
                    if node.left is None:
                        node.left = self._Node()
                    node = node.left
                node.num += 1

        # ---------- 逐个数计算低于low的数量和高于high的数量 ----------
        size = len(nums)
        ans = 0
        for i in range(size):
            num = nums[i]

            # 计算低于low的XOR数量
            min_num = 0
            node = tree
            for k in range(max_bit_length - 1, -1, -1):
                # low当前位为1的情况：若另一个数与当前数相同的话XOR当前位0，小于low；不同的话XOR当前位为1，需继续比较
                if (low >> k) & 1 == 1:
                    if (num >> k) & 1:  # 当前位为1
                        min_num += node.right_num
                        if node.left is None:
                            break
                        node = node.left
                    else:  # 当前位为0
                        min_num += node.left_num
                        if node.right is None:
                            break
                        node = node.right

                # low当前位为0的情况：只能选择另一个数与当前数相同的情况，XOR当前位为0，需继续比较
                else:
                    if (num >> k) & 1:  # 当前位为1
                        if node.right is None:
                            break
                        node = node.right
                    else:  # 当前位为0
                        if node.left is None:
                            break
                        node = node.left

            # 计算高于high的XOR数量
            max_num = 0
            node = tree
            for k in range(max_bit_length - 1, -1, -1):
                # high当前位为1的情况：只能选择另一个数与当前数不同的情况，XOR当前位为1，需继续比较
                if (high >> k) & 1 == 1:
                    if (num >> k) & 1:  # 当前位为1
                        if node.left is None:
                            break
                        node = node.left
                    else:  # 当前位为0
                        if node.right is None:
                            break
                        node = node.right

                # high当前位为0的情况：若另一个数与当前数不同的话XOR当前位1，大于high；相同的话XOR当前位为0，需继续比较
                else:
                    if (num >> k) & 1:  # 当前位为1
                        max_num += node.left_num
                        if node.right is None:
                            break
                        node = node.right
                    else:  # 当前位为0
                        max_num += node.right_num
                        if node.left is None:
                            break
                        node = node.left

            # 计算当前数可以组成的数组数量
            # 一个数自己异或自己的情况下，永远小于low的0，会被min_num包含
            ans += size - min_num - max_num

        # 移除重复的情况
        ans //= 2

        return