560. 和为 K 的子数组

题目描述

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

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

解题

简单直接, 但时间复杂度最高 O(n³)

class Solution {func subarraySum(_ nums: [Int], _ k: Int) -> Int {var total = 0, p = 0for index in 0...(nums.count - 1) {p = indexwhile p < nums.count {if nums[index...p].reduce(0, +) == k {total += 1}p += 1}}return total}
}

优化 `nums[index...p].reduce(0, +)` 的时间复杂度为 O(1)

整体时间复杂度优化为 O(n²)

class Solution {func subarraySum(_ nums: [Int], _ k: Int) -> Int {var total = 0, p = 0, sum = 0for index in 0...(nums.count - 1) {p = indexsum = 0while p < nums.count {sum += nums[p]if sum == k {total += 1}p += 1}}return total}
}

前缀和 + 哈希表优化

func subarraySum(_ nums: [Int], _ k: Int) -> Int {var prefixSumCount: [Int: Int] = [0: 1]var prefixSum = 0var count = 0for num in nums {prefixSum += numif let prefixCount = prefixSumCount[prefixSum - k] {count += prefixCount}prefixSumCount[prefixSum, default: 0] += 1}return count
}

知识拓展 - “前缀和”的使用和理解

前缀和（Prefix Sum）是一种常用的算法技巧，主要用于高效地处理一系列连续子数组或区间求和的问题。前缀和通过预处理数组来减少后续查询的时间复杂度，使得处理某些类型的问题变得更加高效。

前缀和的定义

前缀和数组是原始数组的一个扩展数组，其中的每个元素表示原始数组中从第一个元素到当前元素的累积和。具体来说，给定一个数组 nums，其前缀和数组 prefix 定义如下：

[ \text{prefix}[i] = \sum_{j=0}^{i-1} \text{nums}[j] ]

例如，考虑数组 nums = [1, 2, 3, 4]，其前缀和数组 prefix 为：

[ \text{prefix} = [0, 1, 3, 6, 10] ]

前缀和的构建

构建前缀和数组的时间复杂度为 (O(n))，其中 (n) 是原始数组的长度。下面是一个简单的 Swift 实现：

func buildPrefixSum(nums: [Int]) -> [Int] {var prefix = [0]for num in nums {prefix.append(prefix.last! + num)}return prefix
}let nums = [1, 2, 3, 4]
let prefix = buildPrefixSum(nums: nums)
print(prefix)  // 输出: [0, 1, 3, 6, 10]

前缀和的应用

前缀和可以用于快速计算任意子数组的和。假设我们要计算数组 nums 中从索引 i 到索引 j 的子数组的和，可以通过以下公式实现：

[ \text{sum}(i, j) = \text{prefix}[j+1] - \text{prefix}[i] ]

例如，要计算 nums = [1, 2, 3, 4] 中从索引 1 到索引 3 的子数组 [2, 3, 4] 的和，可以通过 prefix[4] - prefix[1] 得到结果 9。

func subarraySum(prefix: [Int], i: Int, j: Int) -> Int {return prefix[j + 1] - prefix[i]
}let sum = subarraySum(prefix: prefix, i: 1, j: 3)
print(sum)  // 输出: 9

示例问题

示例 1：求一个数组中所有子数组的和

给定一个数组，求出所有子数组的和。利用前缀和可以高效地解决这个问题。

func allSubarraySums(nums: [Int]) -> [Int] {let prefix = buildPrefixSum(nums: nums)var result = [Int]()for i in 0..<nums.count {for j in i..<nums.count {result.append(subarraySum(prefix: prefix, i: i, j: j))}}return result
}let sums = allSubarraySums(nums: nums)
print(sums)  // 输出: [1, 3, 6, 10, 2, 5, 9, 3, 7, 4]

示例 2：连续子数组的最大和

可以利用前缀和数组来求解最大子数组和问题（Kadane’s Algorithm 的另一种实现方式）。

func maxSubArraySum(nums: [Int]) -> Int {let prefix = buildPrefixSum(nums: nums)var maxSum = Int.minvar minPrefix = 0for i in 1..<prefix.count {maxSum = max(maxSum, prefix[i] - minPrefix)minPrefix = min(minPrefix, prefix[i])}return maxSum
}let maxSum = maxSubArraySum(nums: nums)
print(maxSum)  // 输出: 10

小结

前缀和是一种强大的工具，特别适用于处理涉及区间和的算法问题。通过预处理数组，前缀和可以将许多问题的时间复杂度从 (O(n^2)) 降低到 (O(n))，大大提高了效率。在实际应用中，前缀和常用于解决诸如区间求和、连续子数组问题和二维矩阵问题等。

重点说明 - 前缀和数组 prefix 初始化为 [0] 的重要性

前缀和数组 prefix 初始化为 [0] 是非常重要的，这一点在前缀和算法中起到了关键作用。以下是具体原因：

1. 边界条件处理

初始化 prefix 为 [0]，能够很好地处理前缀和的边界条件。例如，当计算从数组起点到某一位置的和时，这样的初始化可以确保正确计算。假设我们需要计算数组 nums 从索引 0 到 i 的和：

[ \text{prefix}[i+1] = \sum_{j=0}^{i} \text{nums}[j] ]

若不初始化为 [0]，在计算 prefix 时将出现索引越界或错误的累计结果。

2. 子数组和的计算

使用前缀和数组时，通常通过以下公式计算任意子数组的和：

[ \text{sum}(i, j) = \text{prefix}[j+1] - \text{prefix}[i] ]

在这种情况下，prefix 数组的第一个元素为 0 可以确保在计算子数组和时不丢失任何信息，并且可以处理 i=0 的情况，即计算从数组开始到 j 的子数组的和。

例如：

let nums = [1, 2, 3, 4]
let prefix = [0, 1, 3, 6, 10]func subarraySum(prefix: [Int], i: Int, j: Int) -> Int {return prefix[j + 1] - prefix[i]
}// 计算从索引0到3的子数组和：[1, 2, 3, 4]
let sum = subarraySum(prefix: prefix, i: 0, j: 3)
print(sum)  // 输出: 10

若未初始化 prefix 为 [0]，则上述公式在计算从索引 0 开始的子数组和时会出错。

3. 代码简洁性和一致性

初始化为 [0] 使得前缀和数组的构建和使用更加简洁和一致，不需要在构建和查询过程中做额外的边界处理。这有助于避免边界条件相关的错误，并使代码更易读。

例如：

func buildPrefixSum(nums: [Int]) -> [Int] {var prefix = [0]for num in nums {prefix.append(prefix.last! + num)}return prefix
}let nums = [1, 2, 3, 4]
let prefix = buildPrefixSum(nums: nums)
print(prefix)  // 输出: [0, 1, 3, 6, 10]