Java 前缀和算法实战：从一维到二维的解题思路

说明：这里使用 long 防止求和溢出。时间复杂度优化为 O(n + m)，空间复杂度 O(n)。

二维前缀和进阶

当数据变成矩阵时，我们需要计算子矩阵的元素和。原理与一维类似，但需要用到容斥原理。

定义 dp[i][j] 为从 (1, 1) 到 (i, j) 的矩形区域和。计算公式为： dp[i][j] = dp[i-1][j] + dp[i][j-1] - dp[i-1][j-1] + arr[i][j]

查询矩形 (x1, y1) 到 (x2, y2) 的和时： sum = dp[x2][y2] - dp[x1-1][y2] - dp[x2][y1-1] + dp[x1-1][y1-1]

import java.util.Scanner;

public class Main {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int m = in.nextInt();
        int q = in.nextInt();

        int[][] arr = new int[n + 1][m + 1];
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= m; j++) {
                arr[i][j] = in.nextInt();
            }
        }

        long[][] dp = new long[n + 1][m + 1];
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= m; j++) {
                dp[i][j] = dp[i - 1][j] + dp[i][j - 1] + arr[i][j] - dp[i - 1][j - 1];
            }
        }

        while (q > 0) {
            int x1 = in.nextInt();
            int y1 = in.nextInt();
            int x2 = in.nextInt();
            int y2 = in.nextInt();
            System.out.println(dp[x2][y2] - dp[x1 - 1][y2] - dp[x2][y1 - 1] + dp[x1 - 1][y1 - 1]);
            q--;
        }
    }
}

寻找数组的中心下标

这道题要求找到一个下标，使得其左侧元素之和等于右侧元素之和。利用前缀和可以快速计算左右两侧的和。

class Solution {
    public int pivotIndex(int[] nums) {
        int n = nums.length;
        int[] f = new int[n];
        int[] g = new int[n];

        // 计算前缀和
        for (int i = 1; i < n; i++) {
            f[i] = f[i - 1] + nums[i - 1];
        }
        // 计算后缀和
        for (int i = n - 2; i >= 0; i--) {
            g[i] = g[i + 1] + nums[i + 1];
        }

        // 查找平衡点
        for (int i = 0; i < n; i++) {
            if (f[i] == g[i]) {
                return i;
            }
        }
        return -1;
    }
}

除自身以外数组的乘积

返回一个数组，其中每个位置的值是原数组除该位置外所有元素的乘积。这其实是前缀积和后缀积的结合。

class Solution {
    public int[] productExceptSelf(int[] nums) {
        int n = nums.length;
        int[] f = new int[n];
        int[] g = new int[n];
        int[] ret = new int[n];

        f[0] = 1;
        g[n - 1] = 1;

        // 计算前缀积
        for (int i = 1; i < n; i++) {
            f[i] = f[i - 1] * nums[i - 1];
        }
        // 计算后缀积
        for (int i = n - 2; i >= 0; i--) {
            g[i] = g[i + 1] * nums[i + 1];
        }

        // 合并结果
        for (int i = 0; i < n; i++) {
            ret[i] = f[i] * g[i];
        }
        return ret;
    }
}

和为 k 的子数组

如果直接枚举子数组，复杂度是 O(n²)。利用前缀和配合哈希表，可以将复杂度降为 O(n)。

核心逻辑是：如果 prefixSum[i] - prefixSum[j] = k，那么 prefixSum[j] = prefixSum[i] - k。我们只需要在哈希表中查找是否存在这个差值。

class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> hash = new HashMap<>();
        hash.put(0, 1); // 初始化，处理从开头开始的子数组
        int sum = 0;
        int ret = 0;

        for (int num : nums) {
            sum += num;
            ret += hash.getOrDefault(sum - k, 0);
            hash.put(sum, hash.getOrDefault(sum, 0) + 1);
        }
        return ret;
    }
}

和可被 K 整除的子数组

这道题与上一题类似，区别在于判断条件变成了整除。利用同余定理：(a - b) % k == 0 等价于 a % k == b % k。

需要注意的是负数取模的结果可能为负，所以公式调整为 (sum % k + k) % k。

class Solution {
    public int subarraysDivByK(int[] nums, int k) {
        Map<Integer, Integer> hash = new HashMap<>();
        hash.put(0 % k, 1);
        int sum = 0;
        int ret = 0;

        for (int num : nums) {
            sum += num;
            int r = (sum % k + k) % k;
            ret += hash.getOrDefault(r, 0);
            hash.put(r, hash.getOrDefault(r, 0) + 1);
        }
        return ret;
    }
}

连续数组

题目要求找到最长的子数组，其中 0 和 1 的数量相等。我们可以将 0 视为 -1，问题就转化为寻找和为 0 的最长子数组。

同样使用前缀和 + 哈希表，但这里存储的是 <前缀和，下标>。为了得到最长长度，如果同一个前缀和出现过，我们保留最早的下标。

class Solution {
    public int findMaxLength(int[] nums) {
        Map<Integer, Integer> hash = new HashMap<>();
        hash.put(0, -1); // 初始状态，前缀和为 0 时下标为 -1
        int len = 0;
        int sum = 0;

        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == 0) {
                sum -= 1;
            } else {
                sum += 1;
            }

            if (hash.containsKey(sum)) {
                len = Math.max(len, i - hash.get(sum));
            } else {
                hash.put(sum, i);
            }
        }
        return len;
    }
}

矩阵区域和

给定一个矩阵和一个整数 k，要求输出一个新矩阵，其中每个位置的值是其周围 k 距离内的元素和。这本质上是二维前缀和的应用，只需注意边界裁剪。

class Solution {
    public int[][] matrixBlockSum(int[][] mat, int k) {
        int m = mat.length;
        int n = mat[0].length;
        int[][] dp = new int[m + 1][n + 1];

        // 构建二维前缀和
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                dp[i][j] = dp[i - 1][j] + dp[i][j - 1] - dp[i - 1][j - 1] + mat[i - 1][j - 1];
            }
        }

        int[][] ret = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                // 计算矩形区域的四个角坐标
                int x1 = Math.max(i - k, 0) + 1;
                int y1 = Math.max(j - k, 0) + 1;
                int x2 = Math.min(i + k, m - 1) + 1;
                int y2 = Math.min(j + k, n - 1) + 1;

                ret[i][j] = dp[x2][y2] - dp[x1 - 1][y2] - dp[x2][y1 - 1] + dp[x1 - 1][y1 - 1];
            }
        }
        return ret;
    }
}

通过以上练习，可以看出前缀和的核心在于'预处理'与'快速查询'。掌握这一模式后，面对复杂的区间统计问题，往往能找到更优的解法。