借这个问题来实现一个「可删除/计数」的 Trie 树

本文正在参加「Java主题月 – Java 刷题打卡」，详情查看活动链接

题目描述

这是 LeetCode 上的 1707. 与数组中元素的最大异或值 ，难度为困难。

Tag : 「Trie」、「字典树」、「二分」

给你一个由非负整数组成的数组 nums 。另有一个查询数组 queries ，其中 queries[i] = [xi, mi] 。

第 i 个查询的答案是 xi 和任何 nums 数组中不超过 mi 的元素按位异或（XOR）得到的最大值。

换句话说，答案是 max(nums[j] XOR xi) ，其中所有 j 均满足 nums[j] <= mi 。如果 nums 中的所有元素都大于 mi，最终答案就是 -1 。

返回一个整数数组 answer 作为查询的答案，其中 answer.length == queries.length 且 answer[i] 是第 i 个查询的答案。

示例 1：

输入：nums = [0,1,2,3,4], queries = [[3,1],[1,3],[5,6]]

输出：[3,3,7]

解释：
1) 0 和 1 是仅有的两个不超过 1 的整数。0 XOR 3 = 3 而 1 XOR 3 = 2 。二者中的更大值是 3 。
2) 1 XOR 2 = 3.
3) 5 XOR 2 = 7.
复制代码

示例 2：

输入：nums = [5,2,4,6,6,3], queries = [[12,4],[8,1],[6,3]]

输出：[15,-1,5]
复制代码

提示：

1 <= nums.length, queries.length <= $10^5$
queries[i].length == 2
0 <= nums[j], xi, mi <= $10^9$

基本分析

在做本题之前，请先确保已经完成 421. 数组中两个数的最大异或值。

这种提前给定了所有询问的题目，我们可以运用离线思想（调整询问的回答顺序）进行求解。

对于本题有两种离线方式可以进行求解。

普通 Trie

第一种方法基本思路是：不一次性地放入所有数，而是每次将需要参与筛选的数字放入 $Trie$ ，再进行与 421. 数组中两个数的最大异或值类似的贪心查找逻辑。

具体的，我们可以按照下面的逻辑进行处理：

对 nums 进行「从小到大」进行排序，对 queries 的第二维进行「从小到大」排序（排序前先将询问原本的下标映射关系存下来）。
按照排序顺序处理所有的 queries[i]：
1. 在回答每个询问前，将小于等于 queries[i][1] 的数值存入 $Trie$ 。由于我们已经事先对 nums 进行排序，因此这个过程只需要维护一个在 nums 上有往右移动的指针即可。
2. 然后利用贪心思路，查询每个 queries[i][0] 所能找到的最大值是多少，计算异或和（此过程与 421. 数组中两个数的最大异或值一致）。
3. 找到当前询问在原询问序列的下标，将答案存入。

代码：

class Solution {
    static int N = (int)1e5 * 32;
    static int[][] trie = new int[N][2];
    static int idx = 0;
    public Solution() {
        for (int i = 0; i <= idx; i++) {
            Arrays.fill(trie[i], 0);
        }
        idx = 0;
    }
    void add(int x) {
        int p = 0;
        for (int i = 31; i >= 0; i--) {
            int u = (x >> i) & 1;
            if (trie[p][u] == 0) trie[p][u] = ++idx;
            p = trie[p][u];
        }
    }
    int getVal(int x) {
        int ans = 0;
        int p = 0;
        for (int i = 31; i >= 0; i--) {
            int a = (x >> i) & 1, b = 1 - a;
            if (trie[p][b] != 0) {
                p = trie[p][b];
                ans = ans | (b << i);
            } else {
                p = trie[p][a];
                ans = ans | (a << i);
            } 
        }
        return ans ^ x;
    }
    public int[] maximizeXor(int[] nums, int[][] qs) {
        int m = nums.length, n = qs.length;

        // 使用哈希表将原本的顺序保存下来
        Map<int[], Integer> map = new HashMap<>();
        for (int i = 0; i < n; i++) map.put(qs[i], i);

        // 将 nums 与 queries[x][1] 进行「从小到大」进行排序
        Arrays.sort(nums);
        Arrays.sort(qs, (a, b)->a[1]-b[1]);

        int[] ans = new int[n];
        int loc = 0; // 记录 nums 中哪些位置之前的数已经放入 Trie
        for (int[] q : qs) {
            int x = q[0], limit = q[1];
            // 将小于等于 limit 的数存入 Trie
            while (loc < m && nums[loc] <= limit) add(nums[loc++]);
            if (loc == 0) {
                ans[map.get(q)] = -1;    
            } else {
                ans[map.get(q)] = getVal(x);    
            }
        }
        return ans;
    }
}
复制代码

时间复杂度：令 nums 的长度为 m，qs 的长度为 n。两者排序的复杂度为 $O(m\log{m})$ 和 $O(n\log{n})$ ；将所有数插入 $Trie$ 和从 $Trie$ 中查找的复杂度均为 $O(Len)$ ， $Len$ 为 $32$ 。

整体复杂度为 $O(m\log{m} + n\log{n} + (m + n) * Len)$ = $O(m * \max(\log{m}, Len) + n * \max(\log{n}, Len))$ 。

空间复杂度： $O(C)$ 。其中 $C$ 为常数，固定为 $1e5 * 32 * 2$ 。

计数 Trie & 二分

另外一个比较「增加难度」的做法是，将整个过程翻转过来：一次性存入所有的 $Trie$ 中，然后每次将不再参与的数从 $Trie$ 中移除。相比于解法一，这就要求我们为 $Trie$ 增加一个「删除/计数」功能，并且需要实现二分来找到移除元素的上界下标是多少。

具体的，我们可以按照下面的逻辑进行处理：

对 nums 进行「从大到小」进行排序，对 queries 的第二维进行「从大到小」排序（排序前先将询问原本的下标映射关系存下来）。
按照排序顺序处理所有的 queries[i]：
1. 在回答每个询问前，通过「二分」找到在 nums 中第一个满足「小于等于 queries[i][1] 的下标在哪」，然后将该下标之前的数从 $Trie$ 中移除。同理，这个过程我们需要使用一个指针来记录上一次删除的下标位置，避免重复删除。
2. 然后利用贪心思路，查询每个 queries[i][0] 所能找到的最大值是多少。注意这是要判断当前节点是否有被计数，如果没有则返回 $-1$ 。
3. 找到当前询问在原询问序列的下标，将答案存入。

代码：

class Solution {
    static int N = (int)1e5 * 32;
    static int[][] trie = new int[N][2];
    static int[] cnt = new int[N];
    static int idx = 0;
    public Solution() {
        for (int i = 0; i <= idx; i++) {
            Arrays.fill(trie[i], 0);
            cnt[i] = 0;
        }
        idx = 0;
    }
    // 往 Trie 存入(v = 1)/删除(v = -1) 某个数 x
    void add(int x, int v) {
        int p = 0;
        for (int i = 31; i >= 0; i--) {
            int u = (x >> i) & 1;
            if (trie[p][u] == 0) trie[p][u] = ++idx;
            p = trie[p][u];
            cnt[p] += v;
        }
    }
    int getVal(int x) {
        int ans = 0;
        int p = 0;
        for (int i = 31; i >= 0; i--) {
            int a = (x >> i) & 1, b = 1 - a;
            if (cnt[trie[p][b]] != 0) {
                p = trie[p][b];
                ans = ans | (b << i);
            } else if (cnt[trie[p][a]] != 0) {
                p = trie[p][a];
                ans = ans | (a << i);
            } else {
                return -1;
            }
        }
        return ans ^ x;
    }
    public int[] maximizeXor(int[] nums, int[][] qs) {
        int n = qs.length;
        
        // 使用哈希表将原本的顺序保存下来
        Map<int[], Integer> map = new HashMap<>();
        for (int i = 0; i < n; i++) map.put(qs[i], i);

        // 对两者排降序
        sort(nums);
        Arrays.sort(qs, (a, b)->b[1]-a[1]);

        // 将所有数存入 Trie
        for (int i : nums) add(i, 1);

        int[] ans = new int[n];
        int left = -1; // 在 nums 中下标「小于等于」left 的值都已经从 Trie 中移除
        for (int[] q : qs) {
            int x = q[0], limit = q[1];
            // 二分查找到待删除元素的右边界，将其右边界之前的所有值从 Trie 中移除。
            int right = getRight(nums, limit);            
            for (int i = left + 1; i < right; i++) add(nums[i], -1);
            left = right - 1;
            ans[map.get(q)] = getVal(x);
        }
        return ans;
    }
    // 二分找到待删除的右边界
    int getRight(int[] nums, int limit) {
        int l = 0, r = nums.length - 1;
        while (l < r) {
            int mid = l + r >> 1;
            if (nums[mid] <= limit) {
                r = mid;
            } else {
                l = mid + 1;
            }
        }
        return nums[r] <= limit ? r : r + 1;
    }
    // 对 nums 进行降序排序（Java 没有 Api 直接支持对基本类型 int 排倒序，其他语言可忽略）
    void sort(int[] nums) {
        Arrays.sort(nums);
        int l = 0, r = nums.length - 1;
        while (l < r) {
            int c = nums[r];
            nums[r--] = nums[l];
            nums[l++] = c;
        }
    }
}
复制代码

时间复杂度：令 nums 的长度为 m，qs 的长度为 n，常数 $Len = 32$ 。两者排序的复杂度为 $O(m\log{m})$ 和 $O(n\log{n})$ ；将所有数插入 $Trie$ 的复杂度为 $O(m * Len)$ ；每个查询都需要经过「二分」找边界，复杂度为 $O(n\log{m})$ ；最坏情况下所有数都会从 $Trie$ 中被标记删除，复杂度为 $O(m * Len)$ 。