本文为Python算法题集之一的代码示例

题目560：和为K的子数组

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

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

示例 1：

输入：nums = [1,1,1], k = 2
输出：2

示例 2：

输入：nums = [1,2,3], k = 3
输出：2

提示：

  • 1 <= nums.length <= 2 * 104
  • -1000 <= nums[i] <= 1000
  • -107 <= k <= 107

- 问题分析

  1. 本题为求子数组和为K，子数组为非空的连续元素
  2. 主要的计算为两个，1是元素求和，而是比较和与K
  3. 基本的遍历为双层循环，从第一个元素开始，计算从此元素开始有多少次和为K，所以基本的时间算法复杂度为（On2）

- 优化思路

  1. 优化的思路，一是减少计算次数，二是减少检索次数和提升检索效率
  2. 标准的优化思路，前缀和【从第1个元素累加到第n个元素的和，称为第n个前缀和】，只要出现2个前缀和的差=k，则这两个下标之间的数组即满足要求
  3. 进一步优化，第n个元素前缀和为isum_n，则1至n-1如果有x个前缀和为k-isum_n，则第n个元素可以在1至n-1之间形成x个和为k的子数组
  4. 以上优化都是考虑前缀和差，因此对于nums只有一个元素的情况要单独进行处理
  1. 标准版【双循环】，毫无疑问，直接超时
    注意：
  • CheckFuncPerf是我写的函数用时和内存占用模块，地址在这里：测量函数运行用时、内存占用的代码单元CheckFuncPerf.py以及使用方法
  • 测试的超长nums文件是官网的，已经上传到CSDN，地址在这里：和为K的子数组测试用超长数组 
import CheckFuncPerf as cfp

def subarraySum_base(nums, k):
    result=0
    for iIdx in range(len(nums)):
        iSum = nums[iIdx]
        if iSum == k:
            result += 1
        for jIdx in range(iIdx+1, len(nums)):
            iSum += nums[jIdx]
            if iSum == k:
                result += 1
    return result

testcase_big1 = open(r'testcase/hot10_big1.txt', mode='r', encoding='utf-8').read().replace('[', '').replace(']', '')
testcase_big1 = testcase_big1.split(',')
nums = [int(x) for x in testcase_big1]
k = 714

result = cfp.getTimeMemoryStr(subarraySum_base, nums, k)
print(result['msg'], '执行结果 = {}'.format(result['result']))
# 运行结果
函数 subarraySum_base 的运行时间为 10591.62 ms；内存使用量为 4.00 KB 执行结果 = 40
  1. 优化版【提前计算前缀和，直接计算前缀和之差】，有想法，然并卵，超时依旧
import CheckFuncPerf as cfp

def subarraySum_ext1(nums, k):
    if len(nums)==1:
        if nums[0] == k:
            return 1
        else:
            return 0
    isumpref=[0] * len(nums)
    iSum = 0
    for iIdx in range(len(nums)):
        iSum += nums[iIdx]
        isumpref[iIdx] = iSum
    result=0
    for iIdx in range(len(nums)):
        if iIdx==0 and nums[iIdx] == k:
            result += 1
        if isumpref[iIdx] == k:
            result += 1
        for jIdx in range(iIdx+1, len(nums)):
            if isumpref[jIdx]-isumpref[iIdx] == k:
                result += 1
    return result

testcase_big1 = open(r'testcase/hot10_big1.txt', mode='r', encoding='utf-8').read().replace('[', '').replace(']', '')
testcase_big1 = testcase_big1.split(',')
nums = [int(x) for x in testcase_big1]
k = 714

result = cfp.getTimeMemoryStr(subarraySum_ext1, nums, k)
print(result['msg'], '执行结果 = {}'.format(result['result']))
# 运行结果
函数 subarraySum_ext1 的运行时间为 10508.87 ms；内存使用量为 732.00 KB 执行结果 = 40
  1. 优化改进版【单层循环，将到第n个元素的所有前缀和的计数存入字典】，优化所见略同，仅超越68%
import CheckFuncPerf as cfp

def subarraySum_ext2(nums, k):
    dictsumpref = {}
    isum, iresult = 0, 0
    for num in nums:
        isum += num
        if isum == k:
            iresult += 1
        if isum - k in dictsumpref:
            iresult += dictsumpref[isum - k]
        dictsumpref[isum] = dictsumpref.get(isum, 0) + 1
    return iresult

testcase_big1 = open(r'testcase/hot10_big1.txt', mode='r', encoding='utf-8').read().replace('[', '').replace(']', '')
testcase_big1 = testcase_big1.split(',')
nums = [int(x) for x in testcase_big1]
k = 714

result = cfp.getTimeMemoryStr(subarraySum_ext2, nums, k)
print(result['msg'], '执行结果 = {}'.format(result['result']))
# 运行结果
函数 subarraySum_ext2 的运行时间为 3.99 ms；内存使用量为 160.00 KB 执行结果 = 40

一日练，一日功，一日不练十日空

may the odds be ever in your favor ~