02、和为K的子数组

厨子大约 7 分钟数据结构算法算法基地面试前缀和哈希表数组子数组和刷题程序厨校招社招

题目描述

leetcode560. 和为 K 的子数组open in new window

给定一个整数数组和一个整数 k,你需要找到该数组中和为 k 的连续的子数组的个数。

示例 1 :

输入:nums = [1,1,1], k = 2 输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。

暴力法

解析

这个题目的题意很容易理解,就是让我们返回和为 k 的子数组的个数,所以我们直接利用双重循环解决该题,这个是很容易想到的。我们直接看代码吧。

代码

Java Code:

class Solution {
    public int subarraySum(int[] nums, int k) {
         int len = nums.length;
         int sum = 0;
         int count = 0;
         for (int i = 0; i < len; ++i) {
             for (int j = i; j < len; ++j) {
                 sum += nums[j];
                 if (sum == k) {
                     count++;
                 }
             }
             sum = 0;
         }
         return count;
    }
}

C++:

#include <vector>  // C++ 中动态数组用 vector,需包含此头文件
using namespace std;  // 简化写法,避免每次写 std::vector

class Solution {
public:
    // C++ 中数组参数用 vector<int>(对应 Java 的 int[]),返回值类型一致
    int subarraySum(vector<int>& nums, int k) {
        int len = nums.size();  // C++ 获取 vector 长度用 size(),对应 Java 的 length
        int sum = 0;
        int count = 0;
        
        // 双层循环逻辑和 Java 完全一致:枚举所有子数组的起始位置 i
        for (int i = 0; i < len; ++i) {
            // 枚举子数组的结束位置 j(从 i 开始,保证子数组连续)
            for (int j = i; j < len; ++j) {
                sum += nums[j];  // 累加子数组元素和
                if (sum == k) {  // 满足目标和 k,计数+1
                    count++;
                }
            }
            sum = 0;  // 切换到下一个起始位置 i,重置累加和
        }
        
        return count;
    }
};

前缀和

解析

下面我们我们使用前缀和的方法来解决这个题目,那么我们先来了解一下前缀和是什么东西。其实这个思想我们很早就接触过了。见下图

我们通过上图发现,我们的 presum 数组中保存的是 nums 元素的和,presum[1] = presum[0] + nums[0];

presum [2] = presum[1] + nums[1],presum[3] = presum[2] + nums[2] ... 所以我们通过前缀和数组可以轻松得到每个区间的和,

例如我们需要获取 nums[2] 到 nums[4] 这个区间的和,我们则完全根据 presum 数组得到,是不是有点和我们之前说的字符串匹配算法中 BM,KMP 中的 next 数组和 suffix 数组作用类似。

那么我们怎么根据 presum 数组获取 nums[2] 到 nums[4] 区间的和呢?见下图

前缀和
前缀和

所以我们 nums[2] 到 nums[4] 区间的和则可以由 presum[5] - presum[2] 得到。

也就是前 5 项的和减去前 2 项的和,得到第 3 项到第 5 项的和。那么我们可以遍历 presum 就能得到和为 K 的子数组的个数啦。

直接上代码。

代码

Java Code:

class Solution {
    public int subarraySum(int[] nums, int k) {
        //前缀和数组
        int[] presum = new int[nums.length+1];
        for (int i = 0; i < nums.length; i++) {
            //这里需要注意,我们的前缀和是presum[1]开始填充的
            presum[i+1] = nums[i] + presum[i];
        }
        //统计个数
        int count = 0;
        for (int i = 0; i < nums.length; ++i) {
            for (int j = i; j < nums.length; ++j) {
                //注意偏移,因为我们的nums[2]到nums[4]等于presum[5]-presum[2]
                //所以这样就可以得到nums[i,j]区间内的和
                if (presum[j+1] - presum[i] == k) {
                    count++;
                }
            }
        }
        return count;
    }
}

C++:

#include <vector>  // C++ 动态数组 vector,对应 Java 的 int[]
using namespace std;

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int n = nums.size();  // C++ 获取数组长度用 size(),对应 Java 的 length
        // 前缀和数组:presum[0] = 0,presum[1] = nums[0],presum[2] = nums[0]+nums[1]...
        vector<int> presum(n + 1, 0);  // 初始化大小为 n+1,默认值 0(Java 是 new int[n+1] 初始为 0)
        
        // 填充前缀和数组(presum[1] 开始对应 nums[0])
        for (int i = 0; i < n; ++i) {
            presum[i + 1] = nums[i] + presum[i];  // 逻辑和 Java 完全一致
        }
        
        int count = 0;
        // 枚举所有子数组 [i, j](i 是起始索引,j 是结束索引,均为 nums 的索引)
        for (int i = 0; i < n; ++i) {
            for (int j = i; j < n; ++j) {
                // 子数组 nums[i..j] 的和 = presum[j+1] - presum[i](和 Java 逻辑一致)
                if (presum[j + 1] - presum[i] == k) {
                    count++;
                }
            }
        }
        return count;
    }
};

前缀和 + HashMap

解析

其实我们在之前的两数之和中已经用到了这个方法,我们一起来回顾两数之和 HashMap 的代码.

Java Code:

class Solution {
    public int[] twoSum(int[] nums, int target) {

        HashMap<Integer,Integer> map  = new HashMap<>();
        //一次遍历
        for (int i = 0; i < nums.length; ++i) {
            //存在时,我们用数组得值为 key,索引为 value
            if (map.containsKey(target - nums[i])){
               return new int[]{i,map.get(target-nums[i])};
            }
            //存入值
            map.put(nums[i],i);
        }
        //返回
        return new int[]{};
    }
}
#include <vector>   // C++ 动态数组 vector,对应 Java 的 int[]
#include <unordered_map>  // C++ 哈希表(无序),对应 Java 的 HashMap
using namespace std;

class Solution {
public:
    // 返回值是 vector<int>(对应 Java 的 int[]),参数 nums 用引用传递(&)提高效率
    vector<int> twoSum(vector<int>& nums, int target) {
        // C++ 中无序哈希表是 unordered_map,键值对类型:<int, int>(对应 Java 的 HashMap<Integer, Integer>)
        unordered_map<int, int> map;
        
        // 一次遍历数组,索引 i 从 0 到 nums.size()-1
        for (int i = 0; i < nums.size(); ++i) {
            // 计算目标补数:target - 当前元素 nums[i]
            int complement = target - nums[i];
            
            // 检查哈希表中是否存在补数(Java 的 map.containsKey() → C++ 的 map.find() != map.end())
            if (map.find(complement) != map.end()) {
                // 存在则返回结果:当前索引 i 和补数对应的索引(map[complement])
                return {i, map[complement]};  // C++ 直接返回初始化的 vector,对应 Java 的 new int[]{...}
            }
            
            // 不存在则将当前元素和索引存入哈希表(Java 的 map.put() → C++ 的 map[] = ...)
            map[nums[i]] = i;
        }
        
        // 题目保证有解,此处返回空数组(对应 Java 的 new int[]{})
        return {};
    }
};

上面代码中,我们将数组的值和索引存入 map 中,当我们遍历到某一值 x 时,判断 map 中是否含有 target - x,即可。其实我们现在这个题目和两数之和原理是一致的,只不过我们是将所有的前缀和前缀和出现的次数存到了 map 里。下面我们来看一下代码的执行过程。

动图解析

代码

class Solution {
    public int subarraySum(int[] nums, int k) {
        if (nums.length == 0) {
            return 0;
        }
        HashMap<Integer,Integer> map = new HashMap<>();
        //细节,这里需要预存前缀和为 0 的情况,会漏掉前几位就满足的情况
        //例如输入[1,1,0],k = 2 如果没有这行代码,则会返回0,漏掉了1+1=2,和1+1+0=2的情况
        //输入:[3,1,1,0] k = 2时则不会漏掉
        //因为presum[3] - presum[0]表示前面 3 位的和,所以需要map.put(0,1),垫下底
        map.put(0, 1);
        int count = 0;
        int presum = 0;
        for (int x : nums) {
            presum += x;
            //当前前缀和已知,判断是否含有 presum - k的前缀和,那么我们就知道某一区间的和为 k 了。
            if (map.containsKey(presum - k)) {
                count += map.get(presum - k);//获取presum-k前缀和出现次数
            }
            //更新
            map.put(presum,map.getOrDefault(presum,0) + 1);
        }
        return count;
    }
}

C++ Code

class Solution
{
public:
    int subarraySum(vector<int> &nums, int k)
    {
        unordered_map<int, int> smp;
        int sum = 0;
        //初始化"最外面"的0
        smp[0] = 1;
        int result = 0;
        for(int i = 0; i < nums.size(); i++)
        {
            sum += nums[i];
            auto mp = smp.find(sum - k);
            if (mp != smp.end())
            {
                //map里面存的一定是在前面的元素
                //可以尝试将map的value换为数组
                result += mp->second;
            }
            smp[sum]++;
        }

        return result;
    }
};