leetcode560：和为 K 的子数组

题目链接

leetcode

题目描述

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

C++ 代码

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

/*
    TLE（Time Limit Exceeded），超时

    以 i 结尾和为 k 的连续子数组个数，统计符合条件下标 j 的个数，
    并且 nums[j..i] 这个子数组的和恰好为 k。
    对于数组中的每个元素 i，枚举所有可能的起始位置 j，判断 [j..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 count = 0;

        for (int i = 0; i < n; i++) {
            int sum = 0;
            for (int j = i; j >= 0; j--) {
                sum += nums[j];
                if (sum == k) {
                    count++;
                }
            }
        }

        return count;
    }
};

/*
    前缀和 + 哈希表

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

    时间复杂度：O(n)
        其中 n 为数组的长度。
        遍历数组的时间复杂度为 O(n)，使用哈希表查询删除的复杂度为 O(1)。
    空间复杂度：O(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; // 初始化

        for (const auto& num : nums) {
            sum += num; // 计算前缀和
            // 如果找到 sum-k，则增加对应的次数
            if (sum2cnt.find(sum-k) != sum2cnt.end()) {
                count += sum2cnt[sum-k];
            }
            sum2cnt[sum]++; // 更新哈希表
        }

        return count;
    }
};

// 辅助函数：打印数组
void printArray(const vector<int>& nums) {
    cout << "[";
    for (size_t i = 0; i < nums.size(); i++) {
        cout << nums[i];
        if (i != nums.size() - 1) cout << ",";
    }
    cout << "]";
}

// 辅助函数：打印二维数组
void printArray(const vector<vector<int>>& nums) {
    cout << "[";
    for (size_t i = 0; i < nums.size(); i++) {
        cout << "[";
        for (size_t j = 0; j < nums[i].size(); j++) {
            cout << nums[i][j];
            if (j != nums[i].size() - 1) cout << ",";
        }
        cout << "]";
        if (i != nums.size() - 1) cout << ",";
    }
    cout << "]";
}

int main() {
    Solution solution;
    vector<vector<int>> nums_cases = { 
        {1,1,1},
        {1,2,3}
    };
    vector<int> k_cases = { 2, 3 };
    
    for (int i = 0; i < nums_cases.size(); i++) {
        auto& nums = nums_cases[i];
        int k = k_cases[i];
        int result = solution.subarraySum(nums, k);
    
        cout << "Input: ";
        printArray(nums);
        cout << ", k = " << k << endl;
        cout << "Output: " << result << endl;
    }
    
    return 0;
}