560 和为K的子数组

一、题目

给你一个整数数组 nums 和一个整数 k ，请你统计并返回 该数组中和为 k 的子数组的个数 。

子数组是数组中元素的连续非空序列。

二、题解

2.1 枚举法

思路： 枚举所有连续子数组。固定子数组的结束位置 i，再让起始位置 j 从 i 往前移动，累加 nums[j] 得到 nums[j..i] 的和 sum，每当 sum == k 就把计数 count 加一。遍历完所有 (i, j) 组合即可统计出全部满足条件的子数组个数。

class Solution {
    public int subarraySum(int[] nums, int k) {
        // count 用来记录和等于 k 的连续子数组数量
        int count = 0;

        // 外层循环：枚举每个子数组的结束位置 i
        for (int i = 0; i < nums.length; i++) {
            // sum 用来记录当前子数组的和
            // 每换一个结束位置 i，都重新从 0 开始累加
            int sum = 0;

            // 内层循环：从 i 往前遍历，枚举所有以 i 结尾的连续子数组
            for (int j = i; j >= 0; j--) {
                // 将 nums[j] 加入当前子数组的和
                // 此时 sum 表示 nums[j] 到 nums[i] 这段连续子数组的和
                sum += nums[j];

                // 如果当前子数组的和等于 k，说明找到一个符合条件的子数组
                if (sum == k) {
                    count++;
                }
            }
        }

        // 返回和等于 k 的连续子数组总数量
        return count;
    }
}

时间复杂度：O(n2)O(n^2)

空间复杂度：O(1)O(1)

2.2 前缀和+哈希表优化

思路： 设前缀和 pre[i] 为 nums[0..i] 之和，则子数组 nums[j+1..i] 的和为 pre[i] - pre[j]。要使其等于 k，即需要存在 pre[j] = pre[i] - k。因此遍历数组、边走边累加前缀和 pre，用哈希表 mp 统计每个前缀和出现的次数；每到一个位置就把 mp 中 pre - k 的出现次数累加到 count，再把当前 pre 计入哈希表。初始放入 mp.put(0, 1) 用于处理「从下标 0 开始」的子数组。

import java.util.HashMap;

class Solution {
    public int subarraySum(int[] nums, int k) {
        // count：记录和等于 k 的连续子数组数量
        // pre：记录从 nums[0] 到当前 nums[i] 的前缀和
        int count = 0, pre = 0;

        // mp 用来记录某个前缀和出现过几次
        // key：前缀和
        // value：这个前缀和出现的次数
        HashMap<Integer, Integer> mp = new HashMap<>();

        // 初始化：前缀和 0 出现 1 次
        // 作用是处理“从下标 0 开始的子数组和正好等于 k”的情况
        mp.put(0, 1);

        // 遍历数组
        for (int i = 0; i < nums.length; i++) {
            // 更新当前前缀和
            // pre 表示 nums[0] + nums[1] + ... + nums[i]
            pre += nums[i];

            // 如果之前出现过前缀和 pre - k
            // 说明从那个位置之后到当前位置 i 的子数组和为 k
            //
            // 因为：
            // 当前前缀和 - 之前某个前缀和 = k
            // pre - (pre - k) = k
            if (mp.containsKey(pre - k)) {
                // 之前有多少个 pre - k，就说明有多少个以 i 结尾的子数组和为 k
                count += mp.get(pre - k);
            }

            // 把当前前缀和 pre 记录到 HashMap 中
            // 如果 pre 之前出现过，就次数 +1
            // 如果 pre 没出现过，就从 0 开始 +1
            mp.put(pre, mp.getOrDefault(pre, 0) + 1);
        }

        // 返回和等于 k 的连续子数组总数量
        return count;
    }
}

时间复杂度：O(n)O(n)

空间复杂度：O(n)O(n)

评论