golang-notes/example/leetcode/018_test.go at main · kesonan/golang-notes · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
//给定一个包含 n 个整数的数组 nums 和一个目标值 target，判断 nums 中是否存在四个元素 a，b，c 和 d ，使得 a + b + c +
// d 的值与 target 相等？找出所有满足条件且不重复的四元组。
//
// 注意：答案中不可以包含重复的四元组。
//
//
//
// 示例 1：
//
//
//输入：nums = [1,0,-1,0,-2,2], target = 0
//输出：[[-2,-1,1,2],[-2,0,0,2],[-1,0,0,1]]
//
//
// 示例 2：
//
//
//输入：nums = [], target = 0
//输出：[]
//
//
//
//
// 提示：
//
//
// 0 <= nums.length <= 200
// -109 <= nums[i] <= 109
// -109 <= target <= 109
//
// Related Topics 数组 哈希表 双指针
// 👍 850 👎 0

package leetcode

import (
	"sort"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestFourSum(t *testing.T) {
	// -2,-1,0,0,1,2
	assert.Equal(t, [][]int{
		{-2, -1, 1, 2},
		{-2, 0, 0, 2},
		{-1, 0, 0, 1},
	}, fourSum([]int{1, 0, -1, 0, -2, 2}, 0))

	assert.Equal(t, [][]int{
		{0, 0, 1, 2},
	}, fourSum([]int{1, 0, -1, 0, -2, 2}, 3))

	assert.Equal(t, [][]int{
		{-2, -1, 1, 2},
		{-1, -1, 1, 1},
	}, fourSum([]int{-2, -1, -1, 1, 1, 2, 2}, 0))

	// -5,-4,-3,-2,1,3,3,5
	assert.Equal(t, [][]int{
		{-5, -4, -3, 1},
	}, fourSum([]int{1, -2, -5, -4, -3, 3, 3, 5}, -11))

	assert.Equal(t, [][]int(nil), fourSum([]int{}, 0))
}

func fourSum(nums []int, target int) [][]int {
	var ret [][]int
	if len(nums) < 4 {
		return ret
	}

	n := len(nums)
	sort.Ints(nums)
	for i, e := range nums {
		if i > 0 && nums[i-1] == e {
			continue
		}

		j := i + 1
		for j < n {
			if j > i+1 && nums[j] == nums[j-1] {
				j += 1
				continue
			}

			var left, right = j + 1, n - 1
			var flag = -2
			for left < right {
				sum := e + nums[j] + nums[left] + nums[right]
				if sum < target {
					left += 1
					flag = -1
				} else if sum > target {
					right -= 1
					flag = 1
				} else {
					switch flag {
					case -1:
						if left > j+1 && nums[left] == nums[left-1] {
							left += 1
							continue
						}
					case 0:
						if left > j+1 && nums[left] == nums[left-1] && right < n-1 && nums[right] == nums[right+1] {
							right -= 1
							continue
						}
					case 1:
						if right < n-1 && nums[right] == nums[right+1] {
							left += 1
							right -= 1
							continue
						}
					}
					ret = append(ret, []int{e, nums[j], nums[left], nums[right]})
					left += 1
					right -= 1
					flag = 0
				}
			}

			j += 1
		}
	}

	return ret
}