来源:力扣(LeetCode)
描述:
给你一个整数数组 nums
(下标 从 0 开始 计数)以及两个整数:low
和 high
,请返回 漂亮数对 的数目。
漂亮数对 是一个形如 (i, j)
的数对,其中 0 <= i < j < nums.length
且 low <= (nums[i] XOR nums[j]) <= high
。
示例 1:
输入:nums = [1,4,2,7], low = 2, high = 6
输出:6
解释:所有漂亮数对 (i, j) 列出如下:
- (0, 1): nums[0] XOR nums[1] = 5
- (0, 2): nums[0] XOR nums[2] = 3
- (0, 3): nums[0] XOR nums[3] = 6
- (1, 2): nums[1] XOR nums[2] = 6
- (1, 3): nums[1] XOR nums[3] = 3
- (2, 3): nums[2] XOR nums[3] = 5
示例 2:
输入:nums = [9,8,4,2,1], low = 5, high = 14
输出:8
解释:所有漂亮数对 (i, j) 列出如下:
- (0, 2): nums[0] XOR nums[2] = 13
- (0, 3): nums[0] XOR nums[3] = 11
- (0, 4): nums[0] XOR nums[4] = 8
- (1, 2): nums[1] XOR nums[2] = 12
- (1, 3): nums[1] XOR nums[3] = 10
- (1, 4): nums[1] XOR nums[4] = 9
- (2, 3): nums[2] XOR nums[3] = 6
- (2, 4): nums[2] XOR nums[4] = 5
提示:
- 1 <= nums.length <= 2 * 104
- 1 <= nums[i] <= 2 * 104
- 1 <= low <= high <= 2 * 104
方法:字典树
思路与算法
题目想要求解有多少对数字的异或运算结果处于 [low, high] 之间,为了方便求解,我们用 f(x) 来表示有多少对数字的异或运算结果小于等于 x,这时问题变为求解 f(high) − f(low − 1)。
考虑枚举一个元素 ai,求解有多少元素 aj (j < i) 使得 ai ⊕ aj ≤ x,其中 ⊕ 表示按位异或运算。由于求解问题时,x 是已知的,我们可以设计一种「从高位到低位依次计算数字个数」的方法,来得到问题的解。该方法的关键点在于:
- 由于数组中的元素都在 [1, 2 × 104 ] 的范围内,那么我们可以将每一个数字表示为一个长度为 15 位的二进制数字(如果不满 15 位,在最高位之前补上若干个前导 0 即可);
- 这 15 个二进制位从低位到高位依次编号为 0, 1, ⋯, 14。我们从最高位第 14 个二进制位开始,依次计算有多少元素与 ai 的异或运算结果小于 x;
- 对于任意一个使得 ai ⊕ aj < x 条件成立的 aj ,都存在一个 k,使得 ai ⊕ aj 的二进制表示中的第 14 位到第 k + 1 位与 x 相同,而第 k 位却小于。
为了更好的计算答案,我们将数组中的元素看做是长度为 15 的字符串,字符串中只包含 0 和 1。如果将字符串放入字典树中,那么在字典树中查询一个字符串的过程,恰好就是从高位开始确定一个二进制的过程。
我们枚举 ai,并将 a0, a1, ⋯, ai-1 作为 aj 放入字典树中,希望找出有多少个 aj 使得 ai ⊕ aj ≤ x。字典树的每个节点记录一个数字,表示有多少个数字以根结点到该节点路径为前缀。为了计算它,我们需要在添加一个数字时,将路径上的所有节点的数字都加 1。
我们可以从字典树的根结点开始遍历,遍历的参照对象是 ai 和 x。假设我们当前遍历到了第 k 个二进制位:
- 如果 x 的第 k 个二进制位为 0,那么此时不存在使得 ai ⊕ aj < x 条件成立的 aj,设 r 是 ai 的第 k 个二进制位,我们需要往表示 r 的子节点走,这保证了路径上的数值与 ai 做异或后前缀与 x 相同。
- 如果 x 的第 k 个二进制位为 1,设 r 是 ai 的第 k 个二进制位,那么此时表示 r 的子节点中记录的数字,就是使得 ai ⊕ aj < x 条件成立的 aj 的个数,将它累加到答案中。然后我们需要往表示 r ⊕ 1 的子节点走,这保证了路径上的数值与 ai 做异或后前缀与 x 相同。
如果在过程中,出现某个子节点不存在使得过程无法继续,我们需要立刻返回答案。否则在最后,我们遍历完所有的 15 个二进制位后,到达的最后一个节点中记录的数字是使得 ai ⊕ aj = x 条件成立的 aj 的个数,也将其累加到答案中。至此,我们求出来所有使得 ai ⊕ aj ≤ x 条件成立的 aj 的个数。
代码:
struct Trie {
// son[0] 表示左子树,son[1] 表示右子树
array<Trie*, 2> son{nullptr, nullptr};
int sum;
Trie():sum(0) {}
};
class Solution {
private:
// 字典树的根节点
Trie* root = nullptr;
// 最高位的二进制位编号为 14
static constexpr int HIGH_BIT = 14;
public:
void add(int num) {
Trie* cur = root;
for (int k = HIGH_BIT; k >= 0; k--) {
int bit = (num >> k) & 1;
if (cur->son[bit] == nullptr) {
cur->son[bit] = new Trie();
}
cur = cur->son[bit];
cur->sum++;
}
}
int get(int num, int x) {
Trie* cur = root;
int sum = 0;
for (int k = HIGH_BIT; k >= 0; k--) {
int r = (num >> k) & 1;
if ((x >> k) & 1) {
if (cur->son[r] != nullptr) {
sum += cur->son[r]->sum;
}
if (cur->son[r ^ 1] == nullptr) {
return sum;
}
cur = cur->son[r ^ 1];
} else {
if (cur->son[r] == nullptr) {
return sum;
}
cur = cur->son[r];
}
}
sum += cur->sum;
return sum;
}
int f(vector<int>& nums, int x) {
root = new Trie();
int res = 0;
for (int i = 1; i < nums.size(); i++) {
add(nums[i - 1]);
res += get(nums[i], x);
}
return res;
}
int countPairs(vector<int>& nums, int low, int high) {
return f(nums, high) - f(nums, low - 1);
}
};
执行用时:248 ms, 在所有 C++ 提交中击败了22.22%的用户
内存消耗:100.7 MB, 在所有 C++ 提交中击败了14.29%的用户
复杂度分析
时间复杂度:O(nlogC)。其中 n 是 nums 的长度,C 是数组中的元素范围。
空间复杂度:O(nlogC)。每一个元素在字典树中需要使用 O(logC) 的空间,因此总空间复杂度为 O(nlogC)。
author:LeetCode-Solution