Skip to content

前缀和与差分数组

一维前缀和与差分数组

前缀和

今天来聊一道简单却十分巧妙的算法问题:算出一共有几个和为 k 的子数组。

图片

最简单的思路就是把所有子数组都穷举出来,算它们的和,看看谁的和等于 k ,但是这种思路的时间复杂度为O(N^2),题目给出的数组长度的数据范围为2*10^4,那么本问题的最终时间复杂度会达到10^8,很显然,这个思路会超时。

关键是, 如何快速得到某个子数组的和呢 ,比如说给你一个数组nums,让你实现一个接口sum(i, j),这个接口要返回nums[i..j]的和,而且会被多次调用,你怎么实现这个接口呢?

因为接口要被多次调用,显然不能每次都去遍历nums[i..j],有没有一种快速的方法在 O(1) 时间内算出nums[i..j]呢?这就需要 前缀和 技巧了。

前缀和的思路是这样的,对于一个给定的数组nums,我们额外开辟一个前缀和数组进行预处理:

int n = nums.length;
// 前缀和数组
int[] preSum = new int[n + 1];
preSum[0] = 0;
for (int i = 0; i < n; i++){
    preSum[i + 1] = preSum[i] + nums[i]; 
}

图片

这个前缀和数组preSum的含义也很好理解,preSum[i]就是nums[0..i-1]的和。那么如果我们想求nums[i..j]的和,只需要一步操作preSum[j+1]-preSum[i]即可,而不需要重新去遍历数组了。

回到这个子数组问题,我们想求有多少个子数组的和为 k,借助前缀和技巧很容易写出一个解法:

    int subarraySum(int[] nums, int k) {
        int n = nums.length;
        // 构造前缀和
        int[] preSum = new int[n + 1];
        for (int i = 0; i < n; i++) {
            preSum[i + 1] = preSum[i] + nums[i];
        }

        int ans = 0;
        // 穷举所有子数组
        for (int i = 1; i <= n; i++) {
            for (int j = 0; j < i; j++) {
                // sum of nums[j..i-1]
                if (preSum[i] - preSum[j] == k) {
                    ans++;
                }
            }
        }

        return ans;
    }

很显然,这种解法虽然利用 前缀和 思想实现了区间和的快速求解,但是总体的时间复杂度还是O(n^2),因为题目给出的数组仍然嵌套循环了2次,依然会超时。 此时,我们可以利用Map结构优化整个数组的遍历过程,使算法整体的时间复杂度将到O(1),具体做法如下:

    int subarraySum(int[] nums, int k) {
        int n = nums.length;

        int ans = 0;
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);

        int sum = 0;
        for (int i = 0; i < n; i++) {
            sum += nums[i];
            int target = sum - k;
            if (map.containsKey(target)) {
                ans += map.get(target);
            }
            map.put(sum, map.getOrDefault(sum, 0) + 1);
        }

        return ans;
    }

可以看到,优化后的思路相当于是使用map数据结构记录前缀和与前缀和出现的次数,每次遍历的时候,将当前数加入到前缀和sum中,然后判断sum-k是否在map中出现过(即判断sum-k在之前的前缀和中是否出现过),将其出现的次数加入到答案即可。

差分数组

前缀和主要适用的场景是原始数组不会被修改的情况下,频繁查询某个区间的累加和 。而 差分数组的主要适用场景是频繁对原始数组的某个区间的元素进行增减 。比如说,给你输入一个数组nums,然后又要求给区间nums[2..6]全部加 1,再给nums[3..9]全部减 3,再给nums[0..4]全部加 2,再给…,一通操作猛如虎,然后问,最后nums数组的值是什么?

常规思路是使用for循环暴力加上或减去val值,这种思路的单次修改时间复杂度是 O(N),对于频繁修改数组的场景效率会很低下。这时就需要用到差分数组的技巧,类似前缀和技巧构造的prefix数组,我们先对nums数组构造一个diff差分数组, diff[i]就是nums[i]nums[i-1]之差

int[] diff = new int[nums.length];
// 构造差分数组
diff[0] = nums[0];
for (int i = 1; i < nums.length; i++) {
    diff[i] = nums[i] - nums[i - 1];
}

图片

通过这个diff差分数组是可以反推出原始数组nums的,代码逻辑如下:

int[] res = new int[diff.length];
// 根据差分数组构造结果数组
res[0] = diff[0];
for (int i = 1; i < diff.length; i++) {
    res[i] = res[i - 1] + diff[i];
}

这样构造差分数组diff,就可以快速进行区间增减的操作 ,如果你想对区间nums[i..j]的元素全部加 3,那么只需要让diff[i] += 3,然后再让diff[j+1] -= 3即可:

图片

原理很简单,回想diff数组反推nums数组的过程,diff[i] += 3意味着给nums[i..]所有的元素都加了 3,然后diff[j+1] -= 3又意味着对于nums[j+1..]所有元素再减 3,那综合起来,就是对nums[i..j]中的所有元素都加 3 了 ,这种思路单次修改数组的时间复杂度为O(1),适用于频繁修改数组的场景。

差分数组修改区间值的代码如下:

public void update(int i, int j, int val) {
    diff[i] += val;
    if (j + 1 < diff.length) {
        diff[j + 1] -= val;
    }
}

注意:j+1 >= diff.length时,说明是对nums[i]及以后的整个数组都进行修改,那么就不需要再给diff数组减val了。

二维前缀和与差分数组

二维前缀和

二维前缀和相较于一维前缀和的整体思路是一样的,不同的是二维前缀和适用于求二维数组的区间和。整个流程包括两个步骤:

步骤一:求 preSum

我们定义preSum[i] [j]表示从[0,0]位置到[i,j]位置的子矩形的所有元素之和。可以用下图帮助我们理解: S(O,D)=S(O,C)+S(O,B)−S(O,A)+D 304.001.jpeg 减去S(O,A)的原因是S(O,C)S(O,B)中都有S(O,A),即加了两次S(O,A),所以需要减去一次S(O,A)。 如果求preSum[i] [j]的话,对应了以下的递推公式: preSum[i] [j] = preSum[i - 1] [j] + preSum[i] [j - 1] - preSum[i - 1] [j - 1] + martrix[i] [j]

步骤二:根据 preSum求子矩形面积 前面已经求出了数组中从[0,0] 位置到 [i,j] 位置的 preSum。下面要利用preSum[i] [j]来快速求出任意子矩形的面积。同样利用一张图来说明: S(A,D)=S(O,D)−S(O,E)−S(O,F)+S(O,G) 304.002.jpeg

加上子矩形S(O,G)面积的原因是S(O,E)S(O,F)中都有S(O,G),即减了两次S(O,G),所以需要加上一次S(O,G)。如果要求[row1,col1][row2,col2]的子矩形的面积的话,用preSum对应了以下的递推公式: preSum[row2][col2]−preSum[row2][col1−1]−preSum[row1−1][col2]+preSum[row1−1][col1−1]

代码

    int[][] sums;

    public NumMatrix(int[][] matrix) {
        int m = matrix.length;
        int n = matrix[0].length;
        sums = new int[m + 1][n + 1];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                sums[i + 1][j + 1] = sums[i][j + 1] + sums[i + 1][j] - sums[i][j] + matrix[i][j];
            }
        }
    }

    public int sumRegion(int row1, int col1, int row2, int col2) {
        return sums[row2 + 1][col2 + 1] - sums[row1][col2 + 1] - sums[row2 + 1][col1] + sums[row1][col1];
    }

二维差分数组

二维差分数组和一维差分数组的思路也是基本相同,对于二维数组matrix,其思路如下:

  1. 如果将矩阵的第(i,j)个单元格中的值增加 1,那么,若对矩阵求二维前缀和,那么下图 (a) 中的黄色区域的值都会增加 1。
  2. 如果要将矩阵中的 任意 区域(如下图中 (b)的蓝色区域)的值增加 1 呢?只需按照下图 (c)来修改矩阵即可。修改后,若对矩阵求前缀和,那么,只会有蓝色的区域的值 +1,其它区域的值都不变。
  3. 最后对差分数组求二维前缀和,就可以求出对应区间的变化量了。

image.png

二维差分数组的代码如下:

    int[][] diff = new int[m + 1][n + 1];

    public void update(int row1, int col1, int row2, int col2) {
        diff[row1][col1] += 1;
        diff[row1][col2 + 1] -= 1;
        diff[row2 + 1][col1] -= 1;
        diff[row2 + 1][col2 + 1] += 1;
    }