【LeetCode】 560. Subarray Sum Equals K

Description

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.

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, 20,000]。
• 陣列中的數字大小範圍為 [-1000, 1000] 且k的大小範圍為 [-1e7, 1e7]。

Example:

Example 1:
Input:nums = [1,1,1], k = 2
Output: 2

Solution

• 很直覺的先想到使用一個sum[i ,j]代表從i加到j的總和，然後使用兩個迴圈跑i和j，找到所有等於k的子陣列。

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

• 觀察一下上面的code，兩層的for都和陣列大小s有關，因此時間複雜度為O(n^2)；送出之後果然吃TLE。

---

• 那麼，我們思考一下如何加速。
• 我們將上面的sum[i, j]改為先得到sum[0, 0] ~ sum[0, s - 1]，也就是先只算原本i = 0的部分。
  • 若輸入為[1, 2, 3]，得到[1, 3, 6]。
• 可以發現sum[i, j]其實就等於sum[0, j] - sum[0, i]。
• 又因為這題沒有要求陣列內容，只要求數量，因此我們可以使用hash跑一次迴圈就完成。
  • hash[x]代表目前所有k + sum[0, i] = x的數量。(也就是說x - sum[0, i]= k)
  • 得到sum[0, i]後，檢查hash[sum[0, i]]是否有值(配對成功)。
  • 更新hash，讓hash[k + sum[0, i]]++。
• 只須注意hash一開始需要讓hash[k] = 1來初始化即可。
• 時間複雜度降為O(n)。

Code

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int s = nums.size();
        int count = 0;
        vector<int> sum_vector(s + 1, 0);
        unordered_map<int, int> hash;
        hash[k] = 1;
        for(int i = 0; i < s; i++)
        {
            sum_vector[i + 1] = sum_vector[i] + nums[i];
            count += hash[sum_vector[i + 1]];
            hash[k + sum_vector[i + 1]] = hash[k + sum_vector[i + 1]] + 1;
        }
        return count;
    }
};

tags: LeetCode C++