算法学习——前缀和

Cumulative Sum

1. 两数之和

func twoSum(nums []int, target int) []int {
    dic := make(map[int]int)
    for i, v := range nums {
        dic[target - v] = i
    }
    for i, v := range nums {
        if k, ok := dic[v]; ok && i != k {
            return []int{i, k}
        }
    }
    return []int{}
}

454. 四数相加 II

转化为两个数组，时间复杂度大大降低，或者直接用字典存起来也可以

func fourSumCount(A []int, B []int, C []int, D []int) int {
    dic := make(map[int]int)
    length := len(A)
    for i := 0; i < length; i++ {
        for j := 0; j < length; j++ {
            dic[-A[i]-B[j]]++
        }
    }
    res := 0
    tmp := 0
    for i := 0; i < length; i++ {
        for j := 0; j < length; j++ {
            tmp = C[i] + D[j]
            if v, ok := dic[tmp]; ok {
                res += v
            }
        }
    }
    return res
}

// 更快，因为这个可能是直接获取了，不用i再去获取？
func fourSumCount(A []int, B []int, C []int, D []int) int {
	mp := make(map[int]int,len(A)*len(B))
	count := 0
	for _, a := range A {
		for _, b := range B {
			mp[a+b]++
		}
	}

	for _, c := range C {
		for _, d := range D {
			if n, ok := mp[-c-d]; ok {
				count += n
			}
		}
	}

	return count
}

560. 和为K的子数组

真正前缀和，通过考虑前面的前缀和当前前缀之间的差值是否等于k来进行判断，如果需要考虑对应区间，可以使用字典存起来对应的起点就好了。

func subarraySum(nums []int, k int) int {
    dic := make(map[int]int)
    dic[0] = 1
    pre, res := 0, 0
    for _, num := range nums {
        pre += num
        if v, ok := dic[pre - k]; ok {
            res += v
        }
        dic[pre]++
    }
    return res
}

1248. 统计「优美子数组」

func numberOfSubarrays(nums []int, k int) int {
    dic := make(map[int]int)
    dic[0] = 1
    cnt, res := 0, 0
    for _, v := range nums {
        if v % 2 == 1 {
            cnt++
        }
        if tmp, ok := dic[cnt - k]; ok {
            res += tmp
        }
        dic[cnt]++
    }
    return res
}