题目
Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.
Example 1: Input:nums = [1,1,1], k = 2 Output: 2
Note: The length of the array is in range [1, 20,000]. The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].
题目的意思就是给一个数组,然后让你找出有多少个连续子数组的和正好是k
分析
那首先最容易想到的就是双重循环遍历每个子序列,然后算子序列的和是否为k,如果满足就加1,这种时间复杂度为O(n³),题目中说到n的范围能到2W,也就是10的4次方,再立方就是12次方了,这显然是会超时的,所以这种暴力方法过不了。
开始也没什么好的思路,看了一下题目的hint,提示可以保存一下两个位置之间的序列和,这样就能避免每次找到一个序列还需要去遍历一遍计算这个序列的和,典型的空间换时间。
具体来说就是用一个二维数组subsum来保存子序列和,subsum[i][j]表示从序列下标i到j之间的子序列的和。这样当求subsum[i][j+1]的时候就只需要做个加法:subsum[i][j+1] = subsum[i][j] + num[j + 1],这种算法把时间复杂度从立方降到了平方,代码如下:
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int ans = 0;
vector<vector<int>> hash(nums.size(), vector<int>(nums.size()));
//长度为1的子序列
for (int i = 0; i < nums.size(); ++i) hash[i][i] = nums[i];
for (int i = 0; i < nums.size(); ++i) {
for (int j = i + 1; j < nums.size(); ++j) {
hash[i][j] = hash[i][j - 1] + nums[j];
if(hash[i][j] == k) ans++;
}
}
return ans;
}
};
这个代码虽然是平方的时间复杂度,但是我觉得过OJ应该没什么问题,但是提交之后发现超了内存…
4次方的数据,二维数组就是16次方,每个数组4个字节…
其实这里是用不到二维数组的,可以压缩成一维,因为我们得到一个子序列和只需要做个判断,所以从第i个位置开始,长度为x的子序列和不影响我们求从i + 1位置开始的长度为x的子序列和,改进之后的代码如下:
解法一:
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int ans = 0;
vector<int>hash(nums.size());
for (int i = 0; i < nums.size(); ++i) {
hash[i] = nums[i];
if (hash[i] == k) ans++;
for (int j = i + 1; j < nums.size(); ++j) {
hash[j] = hash[j - 1] + nums[j];
if (hash[j] == k) ans++;
}
}
return ans;
}
};
终于,这个代码把题目给过了。不出所料,时间性能很差,跑了1500+ms,排在倒数百分之五…不过空间性能还不错,毕竟是O(n)
当然写这篇文章的目的肯定不是介绍这个辣鸡算法,在讨论区学习到一种时间性能在线性的算法,而且我认为这种思想不仅局限在这道题目里的,所以做个学习记录。
按照大神们的称呼,这种方法叫做
prefix sum(前缀数组和)
它的思路就是利用map来记录之前遍历数组时产生的数组的和出现的次数,那么遍历到后面某个位置,如果到这里的数组和与前面某个位置的数组和之差正好是我们想要的值,那我们需要的这个值出现的次数就可以加上之前那个位置值出现的次数。
举个例子:
nums = 3 4 7 2 -3 1 4 2 k = 7
肉眼能看出来的子序列和就是[3,4] [7] [7 2 -3 1] [1 4 2]这四个,我是按照从前往后出现的顺序表示的,这也是遍历的方向。那怎么才能一次遍历就得到这些子序列呢?
先看最后一个[1 4 2]这个子序列,首先要知道这个子序列我们得到的时候是访问到了数组的末尾,所以不妨把它理解成是整个数组的序列和与[3 4 7 2 -3]的差,进一步,我们如果记录了访问到末尾的时候的序列和的值sum,同时也记录了访问到-3的时候[3 4 7 2 -3]的子序列和presum,那么当我们得到sum的时候,用sum - k,得到的值就是presum,我们发现正好是之前出现过的,所以[1 4 2]这部分是可以组成k的子序列。同时我们需要注意出现presum的次数可能不只是1,所以用一个map来记录它出现的次数,这样,我们只需要遍历一次数组,就能从前往后得到所有满足的子序列的数量。这就是前缀数组的思想。来看代码:
解法二:
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
unordered_map<int, int>hash;
int sum = 0, ans = 0;
hash[0] = 1;
for (int i = 0; i < nums.size(); ++i) {
sum += nums[i];
if (hash.count(sum - k)) {
ans += hash[sum - k];
}
hash[sum]++;
}
return ans;
}
};
代码中sum = 0,表示的是从数组开始到当前位置的部分和,ans表示满足的子序列的数量。
然后遍历数组,先更新sum值,当更新到i位置时,如果前面出现了sum - k的值,不妨设这个位置为pre,那么很显然从pos到i的部分子序列和就是k。所以我们只需要把ans加上hash[sum - k]。如果没有出现,那就在hash中记录一下当前位置的sum,这样如果后面某个位置需要sun的时候,就能找到。
注意循环之前把hash0设成了1,原因就是为了解决如果从第一个位置开始的部分子序列满足了的特殊情况,比如上面的例子中,当遍历到第二个4的时候,现在[3 4]满足了,按照程序需要加上hash[7 - 7] = hash[0],所以hash[0]需要设成1。
这种算法思想的时间性能和空间性能都是O(n),所以跑上OJ时间就从1500ms降到了30ms。这种设计的巧妙真心值得好好琢磨和吸收。