leetcode-go/leetcode/medium/18.go

package main

import (
	"sort"
)

func threeSumForTarget(nums []int, target int) [][]int {
	if len(nums) < 3 {
		return [][]int{}
	}
	sort.Ints(nums)
	res := make([][]int, 0)
	for a := 0; a < len(nums)-2; a++ {
		if a > 0 && nums[a] == nums[a-1] {
			continue
		}
		// the same as twoSum: target -> -nums[a], nums -> nums[a+1:]
		// one or more or none solutions
		b, c := a+1, len(nums)-1
		for b < c {
			nb, nc := nums[b], nums[c]
			sum := nums[a] + nb + nc
			if sum < target {
				b++
			} else if sum > target {
				c--
			} else {
				res = append(res, []int{nums[a], nb, nc})
				for ; b < c && nums[b] == nb; b++ {
				}
				for ; b < c && nums[c] == nc; c-- {
				}
			}
		}
	}
	return res
}

func fourSum(nums []int, target int) [][]int {
	sort.Ints(nums)
	res := make([][]int, 0)
	for i := 0; i < len(nums)-3; i++ {
		if i > 0 && nums[i] == nums[i-1] {
			continue
		}
		tmp := threeSumForTarget(nums[i+1:], target-nums[i])
		for _, v := range tmp {
			res = append(res, append([]int{nums[i]}, v...))
		}
	}
	return res
}

/* func main() {
	a1 := []int{1, 0, -1, 0, -2, 2}
	fmt.Println(fourSum(a1, 0))
} */