二维前缀和：高效求解矩阵区域和问题

news2025/2/4 15:46:33

在处理二维矩阵时，频繁计算某一子矩阵的和是一个常见的操作。传统的做法是直接遍历该子矩阵，时间复杂度较高。当矩阵非常大且有大量的查询时，直接计算将变得低效。为了提高效率，我们可以通过 二维前缀和 技巧在常数时间内解决这个问题。

本文将通过一个具体的 Java 实现，介绍如何使用二维前缀和优化子矩阵求和问题。

关键是二维前缀和数组的构造，以及求解区域和的代码部分

测试链接：https://leetcode.cn/problems/range-sum-query-2d-immutable/

一、前缀和的概念

前缀和是解决区间和问题的经典技巧。在一维数组中，前缀和数组 prefixSum 用于存储从数组开头到当前位置的累加和，这样我们可以在 O(1) 时间内查询任意区间 [l, r] 的和。

二维前缀和的思想类似，它在二维矩阵上扩展了前缀和的概念。给定一个 m x n 的矩阵 matrix，二维前缀和数组 sum 中的元素 sum[i][j] 表示从左上角 (0, 0) 到 (i-1, j-1) 的所有矩阵元素的和。通过构造这个前缀和数组，我们能够在常数时间内查询任意子矩阵的元素和。

二、二维前缀和的计算

2.1 二维前缀和的构建

对于一个 m x n 的矩阵 matrix，我们定义一个同样大小的前缀和数组 sum，其中 sum[i][j] 表示从 (0, 0) 到 (i-1, j-1) 的矩阵元素和。构造 sum[i][j] 的公式如下：

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

matrix[i-1][j-1]：当前矩阵元素。
sum[i-1][j]：上方区域的和。
sum[i][j-1]：左侧区域的和。
sum[i-1][j-1]：左上角区域重复计算的部分，需要减去。

这样通过累加计算每个位置的前缀和，最终可以在常数时间内求出任意子矩阵的和。

2.2 子矩阵和的查询

通过上述方式构造的二维前缀和数组，可以快速计算任意子矩阵的元素和。给定一个矩阵区域的左上角 (row1, col1) 和右下角 (row2, col2)，其和可以通过以下公式计算：

sumRegion(row1, col1, row2, col2) = sum[row2+1][col2+1]
                                    - sum[row1][col2+1]
                                    - sum[row2+1][col1]
                                    + sum[row1][col1]

三、Java 实现

以下是使用二维前缀和优化矩阵区域和查询的 Java 实现。我们将使用 NumMatrix 类来实现：

public class NumMatrix {
    private int[][] sum;

    // 构造函数：计算二维前缀和
    public NumMatrix(int[][] matrix) {
        int n = matrix.length;
        int m = matrix[0].length;
        sum = new int[n + 1][m + 1];  // 创建一个多出一行一列的前缀和数组

        // 填充前缀和数组
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= m; j++) {
                sum[i][j] = matrix[i - 1][j - 1] + sum[i - 1][j] + sum[i][j - 1] - sum[i - 1][j - 1];
            }
        }
    }

    // 查询子矩阵的和
    public int sumRegion(int row1, int col1, int row2, int col2) {
        row1++;
        col1++;
        row2++;
        col2++;
        return sum[row2][col2] - sum[row1 - 1][col2] - sum[row2][col1 - 1] + sum[row1 - 1][col1 - 1];
    }

    public static void main(String[] args) {
        // 示例矩阵
        int[][] matrix = {
            {3, 2, 1, 4},
            {1, 5, 3, 2},
            {4, 2, 2, 1},
            {7, 4, 3, 5}
        };

        // 创建 NumMatrix 对象
        NumMatrix numMatrix = new NumMatrix(matrix);

        // 查询子矩阵 (1,1) 到 (2,2) 的和
        System.out.println(numMatrix.sumRegion(1, 1, 2, 2));  // 输出：15
    }
}

3.1 代码分析

构造函数：NumMatrix(int[][] matrix) 用来构造二维前缀和数组 sum。首先，构造一个大小为 (n+1) x (m+1) 的数组，额外的行和列用于处理边界问题。然后通过双重循环填充 sum 数组，利用之前的公式逐步计算前缀和。
sumRegion 方法：sumRegion(int row1, int col1, int row2, int col2) 用于查询子矩阵 (row1, col1) 到 (row2, col2) 的和。通过前缀和的计算公式，能够在常数时间内返回结果。
主函数：在 main 方法中，我们定义了一个 matrix，并创建了 NumMatrix 对象来处理前缀和的计算。然后调用 sumRegion 方法查询从 (1,1) 到 (2,2) 的子矩阵和，输出为 15。