每日算法 — 使用java实现前缀和:一维二维区间查询与差分数组协同

引言

前缀和(Prefix Sum)与差分数组(Difference Array)是处理区间查询和区间更新问题的两把利器。前缀和将区间求和操作转化为两次单点查询的差值,差分数组则将区间批量更新转化为单点更新。本文将用 Java 完整实现这两种技术,并展示它们在一维数组、二维矩阵以及协同使用场景中的强大能力。

前缀和基础

核心思想

对于一个数组 nums,其前缀和数组 prefix 定义为:

prefix[i] = nums[0] + nums[1] + ... + nums[i-1]

特别地,prefix[0] = 0。通过前缀和数组,任意区间 [left, right] 的和可以表示为:

sum(left, right) = prefix[right + 1] - prefix[left]

这样,区间求和的时间复杂度从 O(n) 优化到了 O(1)。

一维前缀和的实现

/**
 * 一维前缀和类
 * 支持 O(1) 时间复杂度的区间求和查询
 */
class PrefixSum1D {
    private long[] prefix;  // 使用 long 防止整数溢出

    /**
     * 构造函数:初始化前缀和数组
     * 时间复杂度:O(n)
     *
     * @param nums 原始数组
     */
    public PrefixSum1D(int[] nums) {
        int n = nums.length;
        prefix = new long[n + 1];  // prefix[0] = 0 作为哨兵

        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + nums[i];
        }
    }

    /**
     * 查询区间 [left, right] 的元素和(包含两端)
     * 时间复杂度:O(1)
     *
     * @param left  区间左端点(包含)
     * @param right 区间右端点(包含)
     * @return 区间和
     */
    public long rangeSum(int left, int right) {
        // 利用前缀和数组的差值快速计算区间和
        // prefix[right+1] 表示 nums[0..right] 的和
        // prefix[left] 表示 nums[0..left-1] 的和
        // 两者相减即为 nums[left..right] 的和
        return prefix[right + 1] - prefix[left];
    }

    /**
     * 查询前缀 [0, index] 的元素和
     */
    public long prefixSum(int index) {
        return prefix[index + 1];
    }
}

应用场景:子数组和等于 K

给定一个整数数组和一个整数 k,找出数组中和为 k 的连续子数组的个数。

import java.util.HashMap;
import java.util.Map;

/**
 * 和为 K 的子数组个数 — 前缀和 + 哈希表
 * 核心思想:如果 prefix[j] - prefix[i] = k,则子数组 [i, j-1] 的和为 k
 * 时间复杂度:O(n)
 * 空间复杂度:O(n)
 *
 * @param nums 整数数组
 * @param k    目标和
 * @return 和为 k 的连续子数组个数
 */
public int subarraySum(int[] nums, int k) {
    // key: 前缀和,value: 该前缀和出现的次数
    Map<Integer, Integer> prefixCount = new HashMap<>();
    prefixCount.put(0, 1);  // 前缀和为0出现1次(空数组)

    int prefixSum = 0;  // 当前前缀和
    int count = 0;      // 满足条件的子数组个数

    for (int num : nums) {
        prefixSum += num;

        // 如果存在 prefixSum - k 的前缀和,说明存在和为 k 的子数组
        // 因为 prefixSum - (prefixSum - k) = k
        count += prefixCount.getOrDefault(prefixSum - k, 0);

        // 记录当前前缀和的出现次数
        prefixCount.put(prefixSum, prefixCount.getOrDefault(prefixSum, 0) + 1);
    }

    return count;
}

二维前缀和

核心思想

对于二维矩阵 matrix,prefix[i][j] 表示从矩阵左上角 (0,0) 到 (i-1,j-1) 的矩形区域内所有元素之和。

二维前缀和的递推公式:

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

子矩阵 [row1,col1] 到 [row2,col2] 的和为:

sum = prefix[row2+1][col2+1] - prefix[row1][col2+1] - prefix[row2+1][col1] + prefix[row1][col1]

二维前缀和的实现

/**
 * 二维前缀和类
 * 支持 O(1) 时间复杂度的子矩阵求和查询
 */
class PrefixSum2D {
    private long[][] prefix;

    /**
     * 构造函数:初始化二维前缀和矩阵
     * 时间复杂度:O(m * n)
     *
     * @param matrix 原始二维矩阵
     */
    public PrefixSum2D(int[][] matrix) {
        if (matrix == null || matrix.length == 0) {
            prefix = new long[1][1];
            return;
        }

        int rows = matrix.length;
        int cols = matrix[0].length;
        prefix = new long[rows + 1][cols + 1];  // 多留一行一列作为哨兵

        for (int i = 1; i <= rows; i++) {
            for (int j = 1; j <= cols; j++) {
                // 容斥原理计算前缀和
                prefix[i][j] = matrix[i - 1][j - 1]
                             + prefix[i - 1][j]      // 上方区域
                             + prefix[i][j - 1]      // 左方区域
                             - prefix[i - 1][j - 1]; // 左上角被重复计算,需要减去
            }
        }
    }

    /**
     * 查询子矩阵 [row1,col1] 到 [row2,col2] 的元素和(包含边界)
     * 时间复杂度:O(1)
     *
     * @param row1 左上角行坐标
     * @param col1 左上角列坐标
     * @param row2 右下角行坐标
     * @param col2 右下角列坐标
     * @return 子矩阵和
     */
    public long regionSum(int row1, int col1, int row2, int col2) {
        // 利用容斥原理,从大的矩形中减去不需要的部分
        return prefix[row2 + 1][col2 + 1]
             - prefix[row1][col2 + 1]     // 减去上方矩形
             - prefix[row2 + 1][col1]     // 减去左方矩形
             + prefix[row1][col1];        // 加回左上角被多减的部分
    }
}

应用场景:最大子矩阵和

/**
 * 最大子矩阵和 — 前缀和 + 动态规划
 * 将二维问题降维到一维,使用 Kadane 算法求解
 * 时间复杂度:O(m² * n) 或 O(m * n²)
 * 空间复杂度:O(n) 或 O(m)
 *
 * @param matrix 二维矩阵
 * @return 最大子矩阵和
 */
public int maxSubmatrixSum(int[][] matrix) {
    if (matrix == null || matrix.length == 0) return 0;

    int rows = matrix.length;
    int cols = matrix[0].length;
    int maxSum = Integer.MIN_VALUE;

    // 枚举上下边界
    for (int top = 0; top < rows; top++) {
        int[] temp = new int[cols];  // temp[j] 表示从 top 行到 bottom 行第 j 列的和

        for (int bottom = top; bottom < rows; bottom++) {
            // 将当前上下边界之间的矩阵压缩为一维数组
            for (int j = 0; j < cols; j++) {
                temp[j] += matrix[bottom][j];
            }

            // 对一维数组使用 Kadane 算法求最大子数组和
            maxSum = Math.max(maxSum, kadane(temp));
        }
    }

    return maxSum;
}

/**
 * Kadane 算法:求一维数组的最大子数组和
 * 时间复杂度:O(n)
 */
private int kadane(int[] nums) {
    int maxEndingHere = nums[0];
    int maxSoFar = nums[0];

    for (int i = 1; i < nums.length; i++) {
        // 以当前元素结尾的最大子数组和
        // 要么单独取当前元素,要么将当前元素接在之前的子数组后面
        maxEndingHere = Math.max(nums[i], maxEndingHere + nums[i]);
        maxSoFar = Math.max(maxSoFar, maxEndingHere);
    }

    return maxSoFar;
}

差分数组

核心思想

差分数组 diff 是前缀和的逆运算。对于原数组 nums,差分数组定义为:

diff[0] = nums[0]
diff[i] = nums[i] - nums[i-1]  (i > 0)

差分数组的核心性质:对原数组的区间 [left, right] 进行批量加值 val,只需修改差分数组的两个位置:

diff[left] += val
diff[right + 1] -= val  (如果 right + 1 在范围内)

差分数组的实现

/**
 * 差分数组类
 * 支持 O(1) 时间复杂度的区间批量更新
 */
class DifferenceArray {
    private int[] diff;    // 差分数组
    private int[] result;  // 结果数组(前缀和还原后的数组)

    /**
     * 构造函数:根据原始数组构建差分数组
     *
     * @param nums 原始数组
     */
    public DifferenceArray(int[] nums) {
        int n = nums.length;
        diff = new int[n];
        diff[0] = nums[0];

        for (int i = 1; i < n; i++) {
            diff[i] = nums[i] - nums[i - 1];
        }
    }

    /**
     * 对区间 [left, right] 内所有元素加上 val
     * 时间复杂度:O(1)
     *
     * @param left  区间左端点(包含)
     * @param right 区间右端点(包含)
     * @param val   要增加的值
     */
    public void rangeAdd(int left, int right, int val) {
        diff[left] += val;
        if (right + 1 < diff.length) {
            diff[right + 1] -= val;
        }
    }

    /**
     * 通过前缀和还原原始数组
     * 时间复杂度:O(n)
     *
     * @return 更新后的数组
     */
    public int[] result() {
        int n = diff.length;
        result = new int[n];
        result[0] = diff[0];

        for (int i = 1; i < n; i++) {
            result[i] = result[i - 1] + diff[i];
        }

        return result;
    }
}

应用场景:航班预订统计

给定一个航班预订表,其中 bookings[i] = [first, last, seats] 表示从 first 到 last 的每个航班都预订了 seats 个座位。返回每个航班的预订座位数。

/**
 * 航班预订统计 — 差分数组经典应用
 * 时间复杂度:O(n),其中 n 为航班数量
 * 空间复杂度:O(n)
 *
 * @param bookings 预订列表,每个元素为 [first, last, seats]
 * @param n        航班总数
 * @return 每个航班的预订座位数
 */
public int[] corpFlightBookings(int[][] bookings, int n) {
    int[] diff = new int[n + 2];  // 多留空间避免越界判断

    for (int[] booking : bookings) {
        int first = booking[0];
        int last = booking[1];
        int seats = booking[2];

        // 区间批量加值:只需修改差分数组的两个端点
        diff[first] += seats;
        diff[last + 1] -= seats;
    }

    // 前缀和还原结果
    int[] result = new int[n];
    int current = 0;
    for (int i = 1; i <= n; i++) {
        current += diff[i];
        result[i - 1] = current;
    }

    return result;
}

前缀和与差分数组的协同使用

在某些复杂场景中,前缀和与差分数组需要协同使用。例如:对数组进行多次区间修改后,再进行多次区间查询。

/**
 * 区间修改 + 区间查询 的综合解决方案
 * 使用"差分的前缀和"或树状数组/线段树
 * 这里展示基于差分数组 + 前缀和的思路
 */
class RangeQueryWithModification {
    private long[] diff;     // 差分数组
    private long[] prefix;   // 差分数组的前缀和(即原数组)
    private long[] prePrefix; // 前缀和的前缀和,用于 O(1) 区间查询

    /**
     * 初始化
     * @param nums 原始数组
     */
    public RangeQueryWithModification(int[] nums) {
        int n = nums.length;
        diff = new long[n + 2];

        // 构建差分数组
        for (int i = 0; i < n; i++) {
            rangeAdd(i, i, nums[i]);  // 将原数组值逐个加入
        }

        buildPrefix(n);
    }

    private void rangeAdd(int left, int right, long val) {
        diff[left + 1] += val;
        diff[right + 2] -= val;
    }

    private void buildPrefix(int n) {
        prefix = new long[n + 1];
        prePrefix = new long[n + 1];

        for (int i = 1; i <= n; i++) {
            prefix[i] = prefix[i - 1] + diff[i];
            prePrefix[i] = prePrefix[i - 1] + prefix[i];
        }
    }

    /**
     * 对区间 [left, right] 进行加值操作后,重新构建前缀和
     */
    public void update(int left, int right, long val) {
        rangeAdd(left, right, val);
        // 实际应用中,多次更新后需要重新构建或使用更高级数据结构
        // 这里简化处理,假设更新后立即重新构建
        buildPrefix(prefix.length - 1);
    }

    /**
     * 查询区间 [left, right] 的和
     */
    public long query(int left, int right) {
        // 利用前缀和的前缀和进行 O(1) 查询
        // sum(left, right) = prePrefix[right+1] - prePrefix[left] - left * prefix[left]
        // 实际实现较为复杂,这里给出简化版本
        long sum = 0;
        for (int i = left; i <= right; i++) {
            sum += prefix[i + 1];
        }
        return sum;
    }
}

完整测试代码

public class PrefixSumDemo {
    public static void main(String[] args) {
        // 一维前缀和测试
        int[] nums1D = {1, 2, 3, 4, 5};
        PrefixSum1D ps1D = new PrefixSum1D(nums1D);
        System.out.println("数组: [1, 2, 3, 4, 5]");
        System.out.println("区间 [1, 3] 的和: " + ps1D.rangeSum(1, 3));  // 2+3+4 = 9
        System.out.println("区间 [0, 4] 的和: " + ps1D.rangeSum(0, 4));  // 15

        // 二维前缀和测试
        int[][] matrix = {
            {1, 2, 3},
            {4, 5, 6},
            {7, 8, 9}
        };
        PrefixSum2D ps2D = new PrefixSum2D(matrix);
        System.out.println("\n子矩阵 [0,0] 到 [1,1] 的和: " + ps2D.regionSum(0, 0, 1, 1));  // 1+2+4+5 = 12
        System.out.println("子矩阵 [1,1] 到 [2,2] 的和: " + ps2D.regionSum(1, 1, 2, 2));  // 5+6+8+9 = 28

        // 差分数组测试
        int[] numsDiff = {0, 0, 0, 0, 0};
        DifferenceArray da = new DifferenceArray(numsDiff);
        da.rangeAdd(1, 3, 5);  // 区间 [1, 3] 加 5
        da.rangeAdd(2, 4, 3);  // 区间 [2, 4] 加 3
        int[] result = da.result();
        System.out.println("\n差分数组更新后: " + java.util.Arrays.toString(result));  // [0, 5, 8, 8, 3]

        // 航班预订测试
        int[][] bookings = {{1, 2, 10}, {2, 3, 20}, {2, 5, 25}};
        int[] flights = corpFlightBookings(bookings, 5);
        System.out.println("航班预订结果: " + java.util.Arrays.toString(flights));  // [10, 55, 45, 25, 25]
    }

    static int[] corpFlightBookings(int[][] bookings, int n) {
        int[] diff = new int[n + 2];
        for (int[] b : bookings) {
            diff[b[0]] += b[2];
            diff[b[1] + 1] -= b[2];
        }
        int[] res = new int[n];
        int cur = 0;
        for (int i = 1; i <= n; i++) {
            cur += diff[i];
            res[i - 1] = cur;
        }
        return res;
    }
}

复杂度分析

技术 预处理时间 单点查询 区间查询 区间更新
普通数组 O(1) O(1) O(n) O(n)
前缀和 O(n) O(1) O(1) O(n)
差分数组 O(n) O(n) O(n) O(1)
前缀和 + 差分 O(n) O(1) O(1) O(1)(均摊)

总结

前缀和与差分数组是算法竞赛和工程实践中最常用的预处理技巧之一。前缀和将区间查询转化为单点查询的差值,差分数组将区间更新转化为单点更新。两者结合可以高效解决”区间修改 + 区间查询”的复合问题。理解这两种技术的本质——前缀和是”累加”,差分数组是”相邻差”——有助于在遇到相关问题时快速识别并应用正确的工具。对于更复杂的动态场景,可以进一步学习树状数组和线段树等进阶数据结构。