给你一个整数数组 nums 和一个整数 k ,请你统计并返回 *该数组中和为 k 的子数组的个数 。
暴力求解
class Solution(object):
def subarraySum(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
count = 0
n = len(nums)
for i in range(n):
for j in range(i,n):
if sum(nums[i:j+1]) == k:
count += 1
return count
时间复杂度为O(n^3)
暴力求解优化
其实,可以对上面方法做一些优化。因为我们求nums[i:j+1]nums[i:j+1]nums[i:j+1]的下一次nums[i:j+1+1]nums[i:j+1+1]nums[i:j+1+1],其实只比前一次多了一个nums[j+1]nums[j+1]nums[j+1]。那么我们就没必要每次都从i开始算了。
class Solution(object):
def subarraySum(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
count = 0
n = len(nums)
for i in range(n):
sum = 0
for j in range(i,n):
sum += nums[j]
if sum == k:
count += 1
return count
时间复杂度为O(n^2)
前缀和
我们通过前缀和数组保存前 n 位的和,presum[1]保存的就是 nums 数组中前 1 位的和,也就是 presum[1] = nums[0], presum[2] = nums[0] + nums[1] = presum[1] + nums[1]. 依次类推,所以我们通过前缀和数组可以轻松得到每个区间的和。
前缀和O(n^2)
class Solution(object):
def subarraySum(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
count = 0
n = len(nums)
preSum = [0]
tmp = 0
for i in range(n):
tmp += nums[i]
preSum.append(tmp)
for i in range(1, n+1):
for j in range(i, n+1):
if preSum[j] - preSum[i-1] == k:
count += 1
return count
preSum[j] 表示的是前 j 个元素的累积和,而 preSum[j] - preSum[i-1] 表示的是索引 i-1 到 j 的子数组的和.
前缀和O(n^2)
class Solution(object):
def subarraySum(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
count = 0
n = len(nums)
preSums = {} # 使用普通字典
preSums[0] = 1
presum = 0
for i in range(n):
presum += nums[i]
# 检查键是否存在,如果存在则使用对应的值,否则假设值为0
if presum - k in preSums:
count += preSums[presum - k]
# 更新前缀和为 presum 的频次
preSums[presum] = preSums.get(presum, 0) + 1
return count
nums = [2, 4, 4, 6, 0, 3, 3]
k = 6
sol = Solution()
result = sol.subarraySum(nums, k)
print(result)