Prefix Sum

总结数组问题

  • 数组不变，区间查询：前缀和、树状数组、线段树；

  • 数组单点修改，区间查询：树状数组、线段树；

  • 数组区间修改，单点查询：差分、线段树；

  • 数组区间修改，区间查询：线段树。

    链接

差分数组

1109. 航班预订统计

方法一：差分数组

航班编号从1开始，初始化diff数组大小为n + 1，diff[0]没有意义

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
    public int[] corpFlightBookings(int[][] bookings, int n) {
        int[] res = new int[n], diff = new int[n + 1];
        for (int[] booking : bookings) {
            diff[booking[0]] += booking[2];
            if (booking[1] != n)
                diff[booking[1] + 1] -= booking[2];
        }
        res[0] = diff[1];
        for (int i = 1; i < n; ++i) 
            res[i] = diff[i + 1] + res[i - 1];
        return res;
    }
}

1094. 拼车

方法一：差分数组

上车位置加，下车位置减

1
2
diff[trip[1]] += trip[0];
diff[trip[2]] -= trip[0];

是diff[trip[2]] -= trip[0]; 而不是diff[trip[2] + 1] -= trip[0];

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution {
    public boolean carPooling(int[][] trips, int capacity) {
        int[] diff = new int[1002];
        for (int[] trip : trips) {
            diff[trip[1]] += trip[0];
            diff[trip[2]] -= trip[0];
        }
        int sum = 0;
        for (int x : diff) {
            sum += x;
            if (sum > capacity)
                return false;
        }
        return true;
    }
}

995. K 连续位的最小翻转次数

方法一：差分数组

题解

遍历到nums[i]时，如果nums[i] + count是偶数，则当前元素(可能被翻转过)实际值是0，需要翻转区间[i, i + k - 1]，

将diff[i + k]减一。遍历到i + k的时候，之前在i时的翻转不再生效，--count(--翻转次数)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution {
    public int minKBitFlips(int[] nums, int k) {
        int n = nums.length;
        int[] diff = new int[n + 1]; // 减少一步越界处理
        int res = 0, count = 0; // 当前被翻转次数
        for (int i = 0; i < n; ++i) {
            if (diff[i] == -1)
                --count;
            if ((nums[i] + count) % 2 == 0) {
                if (i + k > n)
                    return -1;
                ++res;
                ++count;
                --diff[i + k];
            }
        }
        return res;
    }
}

6919. 使数组中的所有元素都等于零

方法一：差分数组

与995. K 连续位的最小翻转次数相似，草稿纸上模拟，多Debug

注意diff[i + k - 1] += nums[i] + curDiff; 复原的位置是i + k - 1

举例

1
2
3
4
5
6
输入：nums = [2,2,3,1,1,0], k = 3
输出：true
解释：可以执行下述操作：
- 选出子数组 [2,2,3] ，执行操作后，数组变为 nums = [1,1,2,1,1,0] 。
- 选出子数组 [2,1,1] ，执行操作后，数组变为 nums = [1,1,1,0,0,0] 。
- 选出子数组 [1,1,1] ，执行操作后，数组变为 nums = [0,0,0,0,0,0] 。
nums 2 2 3 1 1 0 curDiff
diff i = 0 -2 0 2 0 0 0 -2
diff i = 1 -2 0 2 0 0 0 -2
diff i = 2 -2 0 1 0 1 0 -1
diff i = 3 -2 0 1 0 1 0 -1
diff i = 4 -2 0 1 0 1 0 0
diff i = 5 -2 0 1 0 1 0 0 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
    public boolean checkArray(int[] nums, int k) {
        int n = nums.length;
        int[] diff = new int[n];
        int curDiff = 0;
        for (int i = 0; i < n; ++i) {
            if (nums[i] + curDiff > 0) {
                if (i + k > n)
                    return false;
                diff[i] += -nums[i] - curDiff;
                diff[i + k - 1] += nums[i] + curDiff;
            }
            else if (nums[i] + curDiff < 0)
                return false;
            if (diff[i] != 0)
                curDiff += diff[i];
        }
        return true;
    }
}

1893. 检查是否区域内所有整数都被覆盖

方法一：差分数组

1 <= ranges.length <= 50，所以diff大小为52，第0个位置不用，第51个位置用于防止越界

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class Solution {
    public boolean isCovered(int[][] ranges, int left, int right) {
        int[] diff = new int[52];
        int max = 0, min = 50;
        for (int[] range : ranges) {
            ++diff[range[0]];
            --diff[range[1] + 1];
            max = Math.max(max, range[1]);
            min = Math.min(min, range[0]);
        }
        for (int i = min; i <= max; ++i) 
            diff[i] = diff[i - 1] + diff[i];
        for (int i = left; i <= right; ++i)
            if (diff[i] <= 0)
                return false;
        return true;
    }
}

方法二：差分数组 + 前缀和

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution {
    public boolean isCovered(int[][] ranges, int left, int right) {
        int[] diff = new int[52];
        for (int[] range : ranges) {
            ++diff[range[0]];
            --diff[range[1] + 1];
        }
        int preSum = 0;
        for (int i = 1; i <= 50; ++i) {
            preSum += diff[i];
            if (i >= left && i <= right && preSum <= 0)
                return false;
        }
        return true;
    }
}

2251. 花期内花的数目

方法一：差分数组 + TreeMap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
    // O(n),差分数组
    public int[] fullBloomFlowers(int[][] flowers, int[] people) {
        int[] res = new int[people.length];
        TreeMap<Integer, Integer> differMap = new TreeMap<>();
        for (int i = 0; i < flowers.length; i++) {
            differMap.put(flowers[i][0], differMap.getOrDefault(flowers[i][0], 0) + 1);
            differMap.put(flowers[i][1] + 1, differMap.getOrDefault(flowers[i][1] + 1, 0) - 1);
        }
        //差分数组还原
        int pre = 0;
        for (Integer key : differMap.keySet()) {
            pre = differMap.get(key) + pre;
            differMap.put(key, pre);
        }
        for (int i = 0; i < people.length; i++) {
            int x = people[i];
            Integer y = differMap.floorKey(x);
            if (y != null) 
                res[i] = differMap.get(y);
            // res[i] = differMap.floorEntry(people[i]).getValue();
        }
        return res;
    }
}

方法二：二分

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Solution {
    public int[] fullBloomFlowers(int[][] flowers, int[] people) {
        int n = flowers.length;
        int[] start = new int[n], end = new int[n];
        for (int i = 0; i < n; ++i) {
            start[i] = flowers[i][0];
            end[i] = flowers[i][1];
        }
        Arrays.sort(start);
        Arrays.sort(end);
        int[] res = new int[people.length];
        for (int i = 0; i < people.length; ++i) {
            int time = people[i];
            // 时间time开了多少花，枯萎了多少花
            int x = bisearch(start, time), y = bisearch(end, time - 1); // 花在end + 1时枯萎
            res[i] = x - y;
        }
        return res;
    }

    // 找到小于等于t的最大数组下标
    public int bisearch(int[] nums, int t) {
        int l = 0, r = nums.length - 1;
        while (l <= r) {
            int mid = (l + r) >> 1;
            if (nums[mid] > t)
                r = mid - 1;
            else 
                l = mid + 1;
        }
        return l;
    }
}

前缀和

303. 区域和检索 - 数组不可变

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class NumArray {

    int[] preSum;
    int n;

    public NumArray(int[] nums) {
        n = nums.length;
        preSum = new int[n + 1];
        for (int i = 1; i <= n; ++i)
            preSum[i] = preSum[i - 1] + nums[i - 1];
    }
    
    public int sumRange(int left, int right) {
        return preSum[right + 1] - preSum[left];
    }
}

/**
 * Your NumArray object will be instantiated and called as such:
 * NumArray obj = new NumArray(nums);
 * int param_1 = obj.sumRange(left,right);
 */

304. 二维区域和检索 - 矩阵不可变

看图

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class NumMatrix {

    int[][] preSum;
    int m, n;

    public NumMatrix(int[][] matrix) {
        m = matrix.length;
        n = matrix[0].length;
        preSum = new int[m + 1][n + 1];
        for (int i = 1; i <= m; ++i) 
            for (int j = 1; j <= n; ++j) 
                preSum[i][j] = matrix[i - 1][j - 1] - preSum[i - 1][j - 1] + preSum[i][j - 1] + preSum[i - 1][j];
    }
    
    public int sumRegion(int row1, int col1, int row2, int col2) {
        return preSum[row2 + 1][col2 + 1] + preSum[row1][col1] - preSum[row1][col2 + 1] - preSum[row2 + 1][col1];
    }
}

/**
 * Your NumMatrix object will be instantiated and called as such:
 * NumMatrix obj = new NumMatrix(matrix);
 * int param_1 = obj.sumRegion(row1,col1,row2,col2);
 */

363. 矩形区域不超过 K 的最大数值和

题解

方法一：朴素前缀和

1 <= m, n <= 100,时间复杂度\(O(m^2n^2)\),最坏情况\(10^8\)刚刚好。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
    public int maxSumSubmatrix(int[][] matrix, int k) {
        int m = matrix.length, n = matrix[0].length;
        int[][] preSum = new int[m + 1][n + 1];
        for (int i = 1; i <= m; ++i) 
            for (int j = 1; j <= n; ++j) 
                preSum[i][j] = matrix[i - 1][j - 1] + preSum[i][j - 1] + preSum[i - 1][j] - preSum[i - 1][j - 1];
        int max = Integer.MIN_VALUE / 2, sum = 0;
        for (int row1 = 0; row1 < m; ++row1) {
            for (int col1 = 0; col1 < n; ++col1) {
                for (int row2 = row1; row2 < m; ++row2) {
                    for (int col2 = col1; col2 < n; ++col2) {
                        sum = preSum[row2 + 1][col2 + 1] + preSum[row1][col1] - preSum[row1][col2 + 1] - preSum[row2 + 1][col1];
                        if (sum <= k)
                            max = Math.max(max, sum);
                    }
                }
            }
        }
        return max;
    }
}

方法二：前缀和 + 二分查找

题解

523. 连续的子数组和

方法一：前缀和

题解 结论证明

preSum[i]:数组nums从下标0~i-1的和

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution {
    public boolean checkSubarraySum(int[] nums, int k) {
        int n = nums.length;
        int[] preSum = new int[n + 1];
        for (int i = 1; i <= n; ++i)
            preSum[i] = preSum[i - 1] + nums[i - 1];
        Set<Integer> set = new HashSet<>();
        for (int i = 2; i <= n; ++i) {
            set.add(preSum[i - 2] % k);
            if (set.contains(preSum[i] % k))
                return true;
        }
        return false;
    }
}

6952. 统计趣味子数组的数目

方法一：前缀和 + HashMap + 公式转换

类似两数之和

1
2
3
4
5
6
7
8
9
10
11
12
13
public long countInterestingSubarrays(List<Integer> nums, int modulo, int k) {
    Map<Integer, Integer> map = new HashMap<>();
    map.put(0, 1);
    long res = 0;
    int s = 0;
    for (int x : nums) {
        if (x % modulo == k)
            s = (s + 1) % modulo;
        res += map.getOrDefault((s - k + modulo) % modulo, 0);
        map.put(s, map.getOrDefault(s, 0) + 1);
    }
    return res;
}

560. 和为 K 的子数组

方法一：前缀和 + HashMap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
    public int subarraySum(int[] nums, int k) {
        int n = nums.length, res = 0;
        Map<Integer, Integer> cnt = new HashMap<>();
        int s = 0;
        cnt.put(0, 1);
        for (int x : nums) {
            s += x;
            res += cnt.getOrDefault(s - k, 0);
            cnt.put(s, cnt.getOrDefault(s, 0) + 1);
        }
        return res;
    }
}

974. 和可被 K 整除的子数组

方法一：前缀和 + HashMap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
    public int subarraysDivByK(int[] nums, int k) {
        int res = 0, s = 0;
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        for (int x : nums) {
            s += x;
            int mod = (s % k + k) % k;
            res += cnt.getOrDefault(mod, 0);
            cnt.put(mod , cnt.getOrDefault(mod, 0) + 1);
        }
        return res;
    }
}

523. 连续的子数组和

方法一：前缀和 + HashSet

s[r + 1] - s[l] = mk => s[l] % k = s[r + 1] % k

子数组长度至少为2，那么就从2开始遍历

题解

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution {
    public boolean checkSubarraySum(int[] nums, int k) {
        int n = nums.length;
        int[] sum = new int[n + 1];
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + nums[i];
        Set<Integer> set = new HashSet<>();        
        for (int i = 2; i <= n; ++i) {
            set.add(sum[i - 2] % k);
            if (set.contains(sum[i] % k))
                return true;
        }
        return false;
    }
}

525. 连续数组

方法一：前缀和 + HashMap

题解

map.put(0, -1);哨兵，比如可以处理：nums = [0, 1]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution {
    public int findMaxLength(int[] nums) {
        int n = nums.length, res = 0;
        int s = 0;
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, -1);
        for (int i = 0; i < n; ++i) {
            s += (nums[i] == 0 ? -1 : 1);
            if (map.containsKey(s))
                res = Math.max(res, i - map.get(s));
            else
                map.put(s, i);
        }
        return res;
    }
}
Algorithms > Prefix Sum
#前缀和，差分数组
Prefix Sum
https://leopol1d.github.io/2023/07/10/prefix-sum/
作者
Leopold
发布于
2023年7月10日
许可协议