综合笔试题：一道结合了「Top K + 前缀和」的数学题｜Java 刷题打卡

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

题目描述

这是 LeetCode 上的 1738. 找出第 K 大的异或坐标值 ，难度为中等。

Tag : 「Top K」、「数学」、「前缀和」

给你一个二维矩阵 matrix 和一个整数 k ，矩阵大小为 m x n 由非负整数组成。

矩阵中坐标 (a, b) 的值可由对所有满足 0 <= i <= a < m 且 0 <= j <= b < n 的元素 matrix[i][j]（下标从 0 开始计数）执行异或运算得到。

请你找出 matrix 的所有坐标中第 k 大的值（k 的值从 1 开始计数）。

示例 1：

输入：matrix = [[5,2],[1,6]], k = 1

输出：7

解释：坐标 (0,1) 的值是 5 XOR 2 = 7 ，为最大的值。

示例 2：

输入：matrix = [[5,2],[1,6]], k = 2

输出：5

解释：坐标 (0,0) 的值是 5 = 5 ，为第 2 大的值。

示例 3：

输入：matrix = [[5,2],[1,6]], k = 3

输出：4

解释：坐标 (1,0) 的值是 5 XOR 1 = 4 ，为第 3 大的值。

示例 4：

输入：matrix = [[5,2],[1,6]], k = 4

输出：0

解释：坐标 (1,1) 的值是 5 XOR 2 XOR 1 XOR 6 = 0 ，为第 4 大的值。

提示：

m == matrix.length
n == matrix[i].length
1 <= m, n <= 1000
0 <= matrix[i][j] <= $10^6$
1 <= k <= m * n

基本分析

根据题意，我们知道其实就是求「所有子矩阵中第 $k$ 大的异或和」，同时规定所有子矩阵的左上角端点为 $(0, 0)$ 。

数据范围为 $10^3$ ，因此「枚举所有右下角」并「每次计算子矩阵异或和」的朴素做法 $O(m^2 * n^2)$ 不用考虑。

要在全局中找最优，「枚举所有右下角」过程不可避免，我们可以优化「每次计算子矩阵异或和」的过程。

这个分析过程与 1310. 子数组异或查询类似。

异或作为不进位加法，可以利用「偶数次异或结果为 $0$ 」的特性实现类似「前缀和」的容斥。这使得我们可以在 $O(1)$ 的复杂度内计算「某个子矩阵的异或和」。

二维前缀异或 & 优先队列（堆）

创建二维数组 $sum[][]$ ，令 $sum[i][j]$ 为以 $(i, j)$ 为右下角的子矩阵的异或和，我们可以得出计算公式：

sum[i][j] = sum[i - 1][j] ⊕ sum[i][j - 1] ⊕ sum[i - 1][j - 1] ⊕ matrix[i - 1][j - 1]

剩下的问题是，如果从所有的「子矩阵异或和」找到第 $k$ 大的值。

变成了 Top K 问题，可以使用「排序」或「堆」进行求解。

具体的，我们可以建立一个大小为 $k$ 的小根堆，在计算二维前缀异或时，判断当前「子矩阵异或和」是否大于堆顶元素：

大于堆顶元素：当前子矩阵异或和可能是第 $k$ 大的值，堆顶元素不可能为第 $k$ 大的值。将堆顶元素弹出，并将当前子矩阵和加入堆中
小于堆顶元素：不会是第 $k$ 大的值，直接丢弃。
等于堆顶元素：有相同元素在堆中，直接丢弃。

最终的堆顶元素即为答案。

代码：

class Solution {
    public int kthLargestValue(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        int[][] sum = new int[m + 1][n + 1];
        PriorityQueue<Integer> q = new PriorityQueue<>(k, (a, b)->a - b);
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                sum[i][j] = sum[i - 1][j] ^ sum[i][j - 1] ^ sum[i - 1][j - 1] ^ mat[i - 1][j - 1];
                if (q.size() < k) {
                    q.add(sum[i][j]);
                } else {
                    if (sum[i][j] > q.peek()) {
                        q.poll();
                        q.add(sum[i][j]);
                    }
                }
            }
        }
        return q.peek();
    }
}

时间复杂度： $O(m * n * \log{k})$
空间复杂度： $O(m * n)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.1738 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先将所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：github.com/SharingSour…

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。