前缀和算法详解

题目描述:

给定一个整数数组 Array，请计算该数组在每个指定区间内元素的总和。

输入:

第一行输入为整数数组 Array 的长度 n，接下来 n 行，每行一个整数，表示数组的元素。随后的输入为需要计算总和的区间。

5

1

2

3

4

5

0,1

2,4

思路:

对于本题,最直观的解法就是遍历一遍区间,将区间内的数全部相加

cpp 复制代码

#include <iostream>
#include <vector>
using namespace std;
int main() {
    int n, a, b;
    cin >> n;
    vector<int> vec(n);
    for (int i = 0; i < n; i++){
        cin >> vec[i];
    }
    while (cin >> a >> b) {
        int sum = 0;
        for (int i = a; i <= b; i++){
            sum += vec[i];
        }
        cout << sum << endl;
    }
}

但是这样做会超时 ,因为如果我们查询m次 ,每次给的区间都是0~n-1 ,这样算法的时间复杂度就是O(m+n)

对于数据量和查询量大的测试样例,时间花销是非常大的

下面我们引入一个前缀和的概念

前缀和的思想是重复利用计算过的子数组之和，从而降低区间查询需要累加计算的次数

我们使用一个新的数组p,用p[i]去记录原数组中0~i的值之和

i : 0 1 2 3 4 5

vec[] : 1 2 4 6 2 5

p[i] : 1 3 7 13 15 20

这样我们就得到了一个新的数组p[i]

此时

p[5] = vec[0]+vec[1]+vec[2]+vec[3]+vec[4]+vec[5]

p[1] = vec[0]+vec[1]

p[5]-p[1] = vec[2]+vec[3]+vec[4]+vec[5]

我们发现p[5]-p[1]的值就是下标2~5的区间和

那得到了这个规律,我们要求区间和,就只需要拿p数组中的两个元素相减即可

而 p 数组是我们之前就计算好的累加和，所以后面每次求区间和的之后我们只需要 O(1) 的时间操作

cpp 复制代码

#include <iostream>
#include <vector>
using namespace std;
int main() {
    int n;
    cin >> n;
    vector<int> vec(n);
    vector<int> p(n);
    int presum = 0;
    for (int i = 0; i < n; i++) {
        cin >> vec[i];
        presum += vec[i];
        p[i] = presum;
    }
    int a;
    int b;
    while (cin >> a >> b) {
        int sum;
        if (a == 0){
            sum = p[b]
        } ;
        else{
            sum = p[b] - p[a - 1]
        };
        cout << sum << endl;
    }
}

注意:

在求区间时,要注意求解区间,如果区间是2~5,那么使用的是p[5]-p[1]

而不是p[5]-p[2]

我们再来看几道用前缀和求解的题目

【题目描述】:

由于城市规划的限制，只允许将区域按横向或纵向划分成两个子区域，而且每个子区域都必须包含一个或多个区块。

为了确保公平竞争，你需要找到一种分配方式，使得 A 公司和 B 公司各自的子区域内的土地总价值之差最小。

【输入描述】

第一行输入两个正整数，代表 n 和 m。

接下来的 n 行，每行输出 m 个正整数。

【输出描述】

请输出一个整数，代表两个子区域内土地总价值之间的最小差距。

【输入示例】

3 3 1 2 3 2 1 3 1 2 3

【输出示例】

【提示信息】

如果将区域按照如下方式划分：

1 2 | 3 2 1 | 3 1 2 | 3

两个子区域内土地总价值之间的最小差距可以达到 0。

思路:

这题也可以使用前缀和的解法,记录整个二维数组的数据总和sum

先创建一个新数组arr,用arr[i]记录第i行的总和.

再用一个变量count,依次累加arr[i],每次累加完,A公司分得的土地是count,B公司是sum-count,计算土地之差,找到最小值

再重复一遍,计算按列分配的情况

cpp 复制代码

#include<iostream>
#include<vector>
#include<cmath>
#include<algorithm>
#include <climits>
using namespace std;

int main() {
    int n, m;
    cin >> n >> m;
    int sum = 0;
    vector<vector<int>>vec(n, vector<int>(m));
    for (int i = 0;i < n;i++) {
        for (int j = 0;j < m;j++) {
            cin >> vec[i][j];
            sum += vec[i][j];
        }
    }
    
    //按行分配
    vector<int>arr1(n,0);
    for (int i = 0;i < n;i++) {
        for (int j = 0;j < m;j++) {
            arr1[i] += vec[i][j];
        }
    }
    //统计前缀和,找最小差值
    int minNum = INT_MAX;
    int x = 0;
    for (int i = 0;i < n-1;i++) {
        x += arr1[i];
        minNum = min(minNum, abs(x - (sum - x)));
    }

    //按列分配
    vector<int>arr2(m, 0);
    for (int i = 0;i < m;i++) {
        for (int j = 0;j < n;j++) {
            arr2[i] += vec[j][i];
        }
    }
    x = 0;
    for (int i = 0;i < m - 1;i++) {
        x += arr2[i];
        minNum = min(minNum, abs(x - (sum - x)));
    }
    cout << minNum;
    return 0;
}

力扣437:路径总和III

给定一个二叉树的根节点 root ，和一个整数 targetSum ，求该二叉树里节点值之和等于 targetSum 的路径的数目。

路径不需要从根节点开始，也不需要在叶子节点结束，但是路径方向必须是向下的（只能从父节点到子节点）。

思路:

可以使用前缀和加哈希表的思路,每遍历到一个节点,就将该节点的值加到sum中

然后查找哈希表中是否存在num = sum-targetSum,因为sum-num = targetNum,如果存在,就说明该路径下有一个子路径和为targetNum

将当前前缀和存入哈希表中

遍历左子树和右子树进行相同操作

左右子树操作完要删除当前前缀和,相当于换一条路径继续查找(回溯的思想)

cpp 复制代码

class Solution {
public:
    void func(TreeNode* root,int targetSum,unordered_map<long long,int>& mp,long long sum,int& count){
        if(root==nullptr)return;
        sum+=root->val;//累加前缀和
        if(mp.find(sum-targetSum)!=mp.end()){//查找是否存在子路径满足条件
            count+=mp[sum-targetSum];
        }
        mp[sum]++;//加入哈希表中
        func(root->left,targetSum,mp,sum,count);//递归左子树
        func(root->right,targetSum,mp,sum,count);//递归右子树
        mp[sum]--;//删除当前前缀和(回溯)
        if(mp[sum]==0) mp.erase(sum);
    }

    int pathSum(TreeNode* root, int targetSum) {
        unordered_map<long long,int>mp;
        mp[0]++;//如果是空路径,前缀和为0
        int count = 0;
        func(root,targetSum,mp,0,count);
        return count;
    }
};