leetcode560：和为 K 的子数组

题目链接

leetcode

题目描述

给你一个整数数组 nums 和一个整数 k ，统计并返回 该数组中和为 k 的子数组的个数 。
子数组是数组中元素的连续非空序列。

C++ 代码

#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;

void printArray(const vector<int>& nums) {
    cout << "[";
    for (const int& num : nums) {
        cout << num << ",";
    }
    cout << "]";
}

/*
    TLE（Time Limit Exceeded），超出时间限制

    以 i 结尾和为 k 的连续子数组个数，统计符合条件的下标 j 的个数，
    其中 0 <= j <= i 且 nums[j..i] 这个子数组的和恰好为 k。

    对于数组中的每个元素 i，需要枚举所有可能的起始位置 j 来判断 [j..i] 这个子数组和是否为 k。
    这种方法存在性能瓶颈，特别是在数组较大时。
    数学证明：这一步是否可以优化呢？
        s[i] 定义为数组 [0..i] 中所有数的和，s[i] 可以由 s[i−1] 推出，
        公式如下：s[i] = s[i−1] + nums[i]
        [j..i] 这个子数组和为 k，这个条件转化为 s[i] - s[j−1] == k
        移项可知，符合条件的下标 j 需要满足 s[j−1] == s[i] − k

    时间复杂度：O(n^2)
        其中 n 为数组的长度。
        枚举子数组开头和结尾需要时间 O(n^2)，求和需要时间 O(1)。
    空间复杂度：O(1)
        只需要常数空间存放若干变量。
*/
class Solution_0 {
public:
    int subarraySum(vector<int>& nums, int k) {
        int n = nums.size();
        int cnt = 0;
        for (int start = 0; start < n; start ++) {
            int sum = 0;
            for (int end = start; end >= 0; end --) {
                sum += nums[end];
                if (sum == k) cnt ++;
            }
        }

        return cnt;
    }
};

/*
    前缀和 + 哈希表

    计算前缀和：
        首先计算数组的前缀和，即从开始到当前位置的所有元素的和。
    使用哈希表：
        使用一个哈希表来存储每个前缀和出现的次数。
    寻找子数组：
        遍历数组，对于每个位置 i，查找 prefixSum[i] - k 在哈希表中的出现次数，
        这表示从某个位置 j 到 i 的子数组和为 k。

    时间复杂度：O(n)
        其中 n 为数组的长度。
        遍历数组的时间复杂度为 O(n)，使用哈希表查询删除的复杂度均为 O(1)。
    空间复杂度：O(n)
        其中 n 为数组的长度。
        在最坏情况下，哈希表可能有 n 个不同的键值，需要 O(n) 的空间复杂度。
*/
class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int count = 0, sum = 0;
        unordered_map<int, int> sum2cnt;
        sum2cnt[0] = 1; // 初始化，前缀和为 0 出现 1 次
        for (const auto& x : nums) {
            sum += x; // 计算前缀和
            if (sum2cnt.find(sum - k) != sum2cnt.end()) {
                count += sum2cnt[sum - k]; // 如果找到 sum - k，则增加对应的次数
            }
            sum2cnt[sum] ++; // 更新哈希表
        }
        
        return count;
    }
};

int main() {
    Solution solution;
    vector<vector<int>> nums = { 
        {1,3,-1,-3,5,3,6,7},
        {1}
    };
    vector<int> k_nums = { 3, 1 };
    
    for (int i = 0; i < nums.size(); i++) {
        cout << "Input: ";
        printArray(nums[i]);
        cout << ", k = " << k_nums[i] << endl;
        int result = solution.subarraySum(nums[i], k_nums[i]);
        cout << "Output: " << result << endl;
    }
    
    return 0;
}

Golang 代码

package main

import "fmt"

/*
   基本思路：
       使用一个前缀和映射来记录前缀和及其出现的次数。
       遍历数组，计算每个位置的前缀和，并在映射中查找目标前缀和（当前前缀和减去 k）的出现次数。
       将目标前缀和的出现次数累加到结果中。
       最后，返回结果即为和为 K 的子数组的个数。

   时间复杂度：O(n)
       其中 n 是输入数组的长度，只需要遍历数组一次。
   空间复杂度：O(n)
       在最坏情况下，映射可能存储所有不同的前缀和。
*/

func subarraySum(nums []int, k int) int {
    // 存储前缀和及其出现次数的映射
    prefixSumCount := make(map[int]int)
    // 初始化
    prefixSumCount[0] = 1
    res, prefixSum := 0, 0
    for _, num := range nums {
        // 计算当前前缀和, 目标前缀和
        prefixSum += num
        target := prefixSum - k
        // 若目标前缀和存在于映射中，则将其出现次数累加到结果
        if count, exists := prefixSumCount[target]; exists {
            res += count
        }
        prefixSumCount[prefixSum] ++
    }

    return res
}

func main() {
    // 测试用例
	testCases := []struct {
		nums     []int
		k        int
		expected int
	}{
		{[]int{1,1,1}, 2, 2},
		{[]int{1,2,3}, 3, 2},
	}

	for i, tc := range testCases {
		result := subarraySum(tc.nums, tc.k)
		fmt.Printf("Test Case %d, Input: nums = %v, k = %d\n", i+1, tc.nums, tc.k)
		if result == tc.expected {
			fmt.Printf("Test Case %d, Output: %d, PASS\n", i+1, result)
		} else {
			fmt.Printf("Test Case %d, Output: %d, FAIL (Expected: %d)\n", i+1, result, tc.expected)
		}
	}
}