18. 四数之和 @Python @Go

原题

https://leetcode.cn/problems/4sum/description/

思路

循环 + 双指针

复杂度

时间：O(n^3)
空间：O(n)

Python代码

class Solution:
    def fourSum(self, nums: list, target: int):
        nums.sort()
        ans = []
        for a in range(len(nums) - 3):
            # 剪枝
            if nums[a] + nums[a+1] + nums[a+2] + nums[a+3] > target:
                break
            # 去重
            if a > 0 and nums[a] == nums[a-1]:
                continue
            for b in range(a+1, len(nums) - 2):
                # 去重
                if b > a + 1 and nums[b] == nums[b-1]:
                    continue
                c = b + 1 
                d = len(nums) - 1
                while c < d:
                    cur_sum = nums[a] + nums[b] + nums[c] + nums[d]
                    if cur_sum == target:
                        ans.append([nums[a], nums[b], nums[c], nums[d]])
                        while c < d and nums[c+1] == nums[c]:
                            c += 1
                        c += 1 
                        while c < d and nums[d-1] == nums[d]:
                            d -= 1 
                        d -= 1
                    elif cur_sum < target:
                        c += 1 
                    elif cur_sum > target:
                        d -= 1 
        return ans 

Go代码

func fourSum(nums []int, target int) [][]int {
	// 排序
	sort.Ints(nums)
	var ans [][]int
	for a := 0; a < len(nums)-3; a++ {
		// 剪枝
		if nums[a]+nums[a+1]+nums[a+2]+nums[a+3] > target {
			break
		}
		// 去重
		if a > 0 && nums[a] == nums[a-1] {
			continue
		}
		for b := a + 1; b < len(nums)-2; b++ {
			// 去重
			if b > a+1 && nums[b] == nums[b-1] {
				continue
			}
			c, d := b+1, len(nums)-1
			for c < d {
				cur_sum := nums[a] + nums[b] + nums[c] + nums[d]
				if cur_sum == target {
					ans = append(ans, []int{nums[a], nums[b], nums[c], nums[d]})
					// 去重
					for c < d && nums[c+1] == nums[c] {
						c++
					}
					c++
					// 去重
					for c < d && nums[d-1] == nums[d] {
						d--
					}
					d--
				} else if cur_sum < target {
					c++
				} else {
					d--
				}
			}
		}
	}
	return ans
}