一、什么是前缀和
前缀和的思路是这样的,对于一个给定的数组 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,借助前缀和技巧很容易写出一个解法:
class Solution { public: int subarraySum(vector<int>& nums, int k) { //先求数组前缀和 vector<int> presum; int sum=0; for(int i=0;i<nums.size();i++){ sum += nums[i]; presum.push_back(sum); } int res=0; //1 3 6 for(int i=0;i<presum.size();i++){ if(presum[i] == k){ res++; } for(int j=i+1;j<presum.size();j++){ if(presum[j]-presum[i] == k){ res++; } } } return res; } };
这个解法的时间复杂度 O(N^2)
空间复杂度 O(N)
,并不是最优的解法。不过通过这个解法理解了前缀和数组的工作原理之后,可以使用一些巧妙的办法把时间复杂度进一步降低。
二、优化解法
前面的解法有嵌套的 for 循环:
第二层 for 循环在干嘛呢?翻译一下就是,在计算,有几个 j
能够使得 sum[i]
和 sum[j]
的差为 k。毎找到一个这样的 j
,就把结果加一。
我们可以把 if 语句里的条件判断移项,这样写:
if (sum[j] == sum[i] - k)
ans++;
优化的思路是:我直接记录下有几个 sum[j]
和 sum[i] - k
相等,直接更新结果,就避免了内层的 for 循环。我们可以用哈希表,在记录前缀和的同时记录该前缀和出现的次数。
class Solution { public: int subarraySum(vector<int>& nums, int k) { //前缀和且优化时间复杂度 //key为前缀和sum[i],value为前缀和sum[i]的个数,即找sum[j]-k在map中的个数 map<int,int> presum; int sum = 0,res = 0; for(int i=0;i<nums.size();i++){ sum+=nums[i]; if(sum == k) res++; res+= presum[sum-k]; presum[sum]++; } return res; } };