LeetCode Subarray Sum Equals K

LeetCode Subarray Sum Equals K

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:

  1. The length of the array is in range [1, 20,000].
  2. The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].

给定一个数组,问有多少个子数组的和等于k。

暴力方法O(n^3)肯定不行,这种连续子数组的问题,肯定会用到数组的前缀和,所以先把前缀和求出来sum[i],这样任意区间[i,j]之间的和都可以用sum[i]-sum[i-1]求到。不过还是需要O(n^2)的时间,代码如下:

class Solution {
public:
	int subarraySum(vector<int>& nums, int k) {
		int n = nums.size();
		if (n == 0)return 0;
		vector<int> sum(n + 1, 0);
		for (int i = 1; i <= n; ++i)sum[i] = sum[i - 1] + nums[i - 1];
		int ans = 0;
		for (int i = 0; i < n; ++i) {
			for (int j = i + 1; j <= n; ++j) {
				if (sum[j] - sum[i] == k)++ans;
			}
		}
		return ans;
	}
};

本代码提交AC,用时533MS。

看了下tag要用map来做。我们把所有前缀和及其频率存到一个map中,假设当前前缀和为sum,则如果sum-k也在map中,说明由产生sum-k前缀和的位置到当前位置的和是k。

比如样例[1,1,1],遇到第一个1时,前缀和为1,之前没有前缀和为1的情况,所以map[1]=1。当加到最后一个数时,前缀和sum=3,sum-k=1也在map中,正好说明从1开始到结束的位置的和为2,如果map[1]=2,说明有两个前缀和等于1的位置,也就有两个连续和为2的子数组。

代码如下,注意只能在遍历的同时求ans,不能等map都构建好了再求ans,这样保证前缀和为sum-k的位置肯定在i前面,也就是我们统一前缀和为k的结束位置为i,做到不重不漏。

class Solution {
public:
	int subarraySum(vector<int>& nums, int k) {
		int n = nums.size();
		if (n == 0)return 0;
		
		int ans = 0, sum = 0;
		unordered_map<int, int> cnt;
		for (int i = 0; i < n; ++i) {
			++cnt[sum];
			sum += nums[i];
			ans += cnt[sum - k];
		}
		return ans;
	}
};

本代码提交AC,用时36MS。

Leave a Reply

Your email address will not be published. Required fields are marked *