前缀和与差分数组 | 程序员小乐

算法前缀和差分

算法

发布日期: 2024-09-22

作者: 小乐

文章字数: 584

阅读时长: 2 分

阅读次数:

一、前缀和

1.1 定义

前缀和数组 prefix[i] 表示原数组从0到i的元素之和。

原数组: [1, 2, 3, 4, 5]
前缀和: [1, 3, 6, 10, 15]

1.2 构建

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

1.3 区间查询

// 查询[left, right]的和
int query(int left, int right) {
    return prefix[right + 1] - prefix[left];
}

时间复杂度：( O(1) )

1.4 二维前缀和

int[][] prefix = new int[m + 1][n + 1];
for (int i = 1; i <= m; i++) {
    for (int j = 1; j <= n; j++) {
        prefix[i][j] = matrix[i-1][j-1] 
                     + prefix[i-1][j] 
                     + prefix[i][j-1] 
                     - prefix[i-1][j-1];
    }
}

// 查询矩形区域和
int query(int r1, int c1, int r2, int c2) {
    return prefix[r2+1][c2+1] - prefix[r1][c2+1] 
         - prefix[r2+1][c1] + prefix[r1][c1];
}

二、差分数组

2.1 定义

差分数组 diff[i] 表示原数组第i个元素与第i-1个元素的差。

原数组: [1, 3, 6, 10, 15]
差分数组: [1, 2, 3, 4, 5]

2.2 构建

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

2.3 区间更新

对 [left, right] 区间加 val：

diff[left] += val;
if (right + 1 < n) diff[right + 1] -= val;

时间复杂度：( O(1) )

2.4 还原原数组

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

三、经典应用

3.1 航班预订统计

public int[] corpFlightBookings(int[][] bookings, int n) {
    int[] diff = new int[n];
    for (int[] booking : bookings) {
        int first = booking[0] - 1, last = booking[1] - 1, seats = booking[2];
        diff[first] += seats;
        if (last + 1 < n) diff[last + 1] -= seats;
    }
    
    for (int i = 1; i < n; i++) {
        diff[i] += diff[i - 1];
    }
    return diff;
}

3.2 和为K的子数组

public int subarraySum(int[] nums, int k) {
    Map<Integer, Integer> prefixCount = new HashMap<>();
    prefixCount.put(0, 1);
    
    int sum = 0, count = 0;
    for (int num : nums) {
        sum += num;
        count += prefixCount.getOrDefault(sum - k, 0);
        prefixCount.put(sum, prefixCount.getOrDefault(sum, 0) + 1);
    }
    return count;
}