LeetCode 560 和为K的子数组

LeetCode 560 和为K的子数组

给一个数组和一个整数k,找到数组中和为k的连续的子数组的个数。

最直观的方法是前缀和。用一个数组cumulative保存从第0个元素到当前元素的和,cumulative数组中第j个元素与第i个元素的差(cumulative[j] - cumulative[i])即是原数组nums中第i个元素到第j个元素的和(nums[i:j]),通过两层循环计算每两个元素间的累积和。时间复杂度是O(n^2),空间复杂度是O(n)。

// C++
class Solution {
public:
    int subarraySum(vector<int> &nums, int k) {
        vector<int> cul(nums.size() + 1, 0);
        int res = 0;
        for (int i = 1; i <= nums.size(); ++i)
            cul[i] = cul[i - 1] + nums[i - 1];
        for (int i = 0; i < cul.size(); ++i)
            for (int j = i + 1; j < cul.size(); ++j)
                res += cul[j] - cul[i] == k;
        return res;
    }
};

当然也可以直接在原数组上进行操作而不是重新申请一个数组来保存前缀和。空间复杂度是O(1)。

// C++
class Solution {
public:
    int subarraySum(vector<int> &nums, int k) {
        int size = nums.size(), res = 0;
        for (int i = 0; i < size; ++i) {
            int sum = 0;
            for (int j = i; j < size; ++j) {
                sum += nums[j];
                res += sum == k ? 1 : 0;
            }
        }
        return res;
    }
};

在此基础上可以先遍历数组nums,计算从第0个元素到当前元素的和,用哈希表保存出现过的累积和sum的次数。如果sum - k在哈希表中出现过,则代表从当前下标i往前有连续的子数组的和为sum。时间复杂度是O(n),空间复杂度是O(1)。

// C++
class Solution {
public:
    int subarraySum(vector<int> &nums, int k) {
        int sum = 0, res = 0;
        unordered_map<int, int> cul;
        cul[0] = 1;
        for (auto &m : nums) {
            sum += m;
            res += cul[sum - k];
            ++cul[sum];
        }
        return res;
    }
}
# Python3
class Solution:
    def subarraySum(self, nums: 'List[int]', k: 'int') -> 'int':
        sum, res, cul = 0, 0, {}
        cul[0] = 1
        for i in range(len(nums)):
            sum += nums[i]
            if sum - k in cul:
                res += cul[sum - k]
            if sum not in cul:
                cul[sum] = 0
            cul[sum] += 1
        return res
comments powered by Disqus