923.三数之和的多种可能

115 阅读1分钟

题目:
给定一个整数数组 arr ,以及一个整数 target 作为目标值,返回满足 i < j < k 且 arr[i] + arr[j] + arr[k] == target 的元组 i, j, k 的数量。

由于结果会非常大,请返回 109 + 7 的模。
算法:
方法一:暴力
arr元素多,且全为0,且target=0时,TLE

func threeSumMulti(arr []int, target int) int {
	n := len(arr)
	mod := 1000000007
	ans := 0 
	for i := 0; i < n; i ++ {
		for j := i + 1; j < n; j ++ {
			for k := j + 1; k < n; k ++ {
				if arr[i] + arr[j] + arr[k] == target {
					ans = (ans + 1) % mod
				}
			}
		}
	}
	return ans
}

方法二:三指针
首先,i < j < k是一个陷阱,它只要求是三个不等的数,不是要求三个升序的数字如arr[i]<arr[j]<arr[k]。 我们2数之和,我们先对arr排序,然后用双指针,伪代码如下: i, j := 0, n - 1 for i < j { if arr[i] + arr[j] < target { // 小了,增大 i ++ } else if arr[i] + arr[j] > target { // 大了,减少 j -- } else { count ++ i ++ j -- } } 我们可以将三指针问题转换为双指针

func threeSumMulti(arr []int, target int) int {
	n := len(arr)
	mod := 1000000007
	ans := 0 
	sort.Ints(arr)
	for i := 0; i < n; i ++ {
		T := target - arr[i]
		if T < 0 {
			continue
		}
		j := i + 1
		k := n - 1
		for j < k {
			if arr[j] + arr[k] < T {
				j ++
			} else if arr[j] + arr[k] > T {
				k --
			} else {
				if arr[j] != arr[k] {
					jCount, kCount := 1, 1
					for j + 1 < k && arr[j] == arr[j + 1] {
						j ++
						jCount ++
					} 
					for j < k - 1 && arr[k - 1] == arr[k] {
						k --
						kCount ++
					}
					ans = (ans + jCount * kCount) % mod
					j ++
					k --
				} else {
					// [j,k]总共x=k - j + 1个数,从中取2个,C(n/2)=n!/2!/(n-2)!
					ans = (ans + (k - j + 1)*(k - j) / 2) % mod
					// arr[j,k]都是相等的数了,break循环
					break
				}
			}
		}
	}
	return ans
}