算法 - 和为 K 的子数组 (Subarray Sum Equals K)

在 Leetcode 上实现这个算法时，总是运行超时，于是深入的学习了下

• 题目详细描述

先上性能测试结果：

Benchmark_numSubarraysWithSum1            20      78277388 ns/op           0 B/op          0 allocs/op
Benchmark_numSubarraysWithSum2          2000        869859 ns/op        8192 B/op          1 allocs/op
Benchmark_numSubarraysWithSum3          5000        298845 ns/op           0 B/op          0 allocs/op
Benchmark_numSubarraysWithSum4         30000         43751 ns/op       40993 B/op          2 allocs/op

四种方法

golang 实现源码及逻辑分析

package main

import (
    "fmt"
    "math/rand"
)

var (
    // SourceArray: 源数组
    SourceArray = make([]int, 1000)

    // K: 目标值
    K = len(SourceArray) / 10
)

func init() {
    // 为了避免文档过大，数组未静态定义，而在此包初始化时生成
    // 随机填充 0、1，作为数组元素
    for i := range SourceArray {
        SourceArray[i] = rand.Intn(2)
    }
}

// 方法#1
// * 时间复杂度: O(n^3)
// 逻辑最直观，但复杂度度最高，性能很差
func numSubarraysWithSum1(SourceArray []int, K int) int {
    var count int
    for i := 0; i < len(SourceArray); i++ {
        for j := i; j < len(SourceArray); j++ {
            sum := 0
            for _, v := range SourceArray[i : j+1] {
                sum += v
            }
            if sum == K {
                count++
            }
        }
    }
    return count
}

// 方法#2
// * 时间复杂度: O(n^2)
// 理解起来相对复杂，但还可以接受，复杂度一般
func numSubarraysWithSum2(SourceArray []int, K int) int {
    count := 0

    // 生成累计数组 (复杂度: O(n))
    // sumArray[0] -> 0
    // sumArray[1] -> 0 + SourceArray[0]
    // sumArray[2] -> 0 + SourceArray[0] + SourceArray[1] = sumArray[1] + SourceArray[1]
    // ...
    // sumArray[n] -> sumArray[n-1] + SourceArray[n]
    sumArray := make([]int, len(SourceArray)+1)
    sumArray[0] = 0
    for i := range SourceArray {
        sumArray[i+1] = sumArray[i] + SourceArray[i]
    }

    // 计算连续累计值间的差值，就等价于子数组元素之和，如果等于目标值target，则为满足条件  (复杂度: O(n^2))
    // sumArray[1] - sumArray[0] -> mathSum(SourceArray[0:1])
    // sumArray[2] - sumArray[0] -> mathSum(SourceArray[0:2])
    // ...
    // sumArray[2] - sumArray[1] -> mathSum(SourceArray[1:2])
    // ...
    // sumArray[n] - sumArray[m] -> mathSum(SourceArray[n:m])
    for i := range SourceArray {
        for j := range SourceArray {
            if sumArray[j+1]-sumArray[i] == K {
                count++
            }
        }
    }

    return count
}

// 方法#3: 遍历过程中直接求和
// * 时间复杂度: O(n^2)
// 是对方法#1的优化，避免了求和那层遍历，比较好理解，性能也还可以
func numSubarraysWithSum3(SourceArray []int, K int) int {
    count := 0
    for i := 0; i < len(SourceArray); i++ {
        sum := 0
        for j := i; j < len(SourceArray); j++ {
            sum += SourceArray[j]
            if sum == K {
                count++
            }
        }
    }
    return count
}

// 方法#4: 最优算法
// * 时间复杂度: O(n)
// 逻辑很难理解，但性能好的一逼，尝试理解下~
// 按着代码的思路，貌似没啥毛病，但是这种计算方式，并不符合人类的正常思维，而是一个抽象的数学公式，还是死记硬背记下它吧
func numSubarraysWithSum4(SourceArray []int, K int) int {
    count, sum := 0, 0

    // sumMap: MAP[连续累计值 -> 这个累计值出现的次数]
    sumMap := make(map[int]int, len(SourceArray))

    // 首先这个定义从人类思维上就没法理解，但后面确实满足了公式计算的需求
    sumMap[0] = 1

    for i := range SourceArray {
        // sum为当前累计值 = 之前元素之和
        // - 如果下个元素非0，则当前累计值将会发生变化，当前累计值的出现的次数只会为1
        // - 如果出现连续的元素为0，sum不会变，在sumMap[sum]++时，累计值出现的次数会大于1
        // 所以：累计值出现的次数 = 1 + 连续出现0的次数，等价于和为当前累计值对于的子数组的数量
        sum += SourceArray[i]

        // sum-K 为距离目标值的差值
        // - 如果=0，代表正好满足，匹配计数加一，所以可以看出开头定义的`sumMap[0] = 1`是为了此处
        // - 如果<0或者超出Map，代表已经超出目标值，不符合条件，从golang数组里取负数索引会得到0
        // - 如果>0且未超出Map，代表需要补充的差值，从sumMap中获得此差值的数量，同时意味着匹配数量
        count += sumMap[sum-K]
        sumMap[sum]++
    }

    return count
}

func main() {
    fmt.Println("#1:", numSubarraysWithSum1(SourceArray, K))
    fmt.Println("#2:", numSubarraysWithSum2(SourceArray, K))
    fmt.Println("#3:", numSubarraysWithSum3(SourceArray, K))
    fmt.Println("#4:", numSubarraysWithSum4(SourceArray, K))
}

参考

• 官方提供的 4 种 JAVA 算法