leetcode面试经典150题- 15. 三数之和

https://leetcode.cn/problems/3sum/description/?envType=study-plan-v2&envId=top-interview-150

package leetcode150

import (
    "sort"
    "testing"
)

func TestThreeSum(t *testing.T) {
    nums := []int{0, 2, 2, 3, 0, 1, 2, 3, -1, -4, 2}
    res := threeSum(nums)
    for _, re := range res {
        for _, i := range re {
            print(i)
            print(",")
        }
        println()
    }
}

func threeSum(nums []int) [][]int {
    if len(nums) < 3 {
        return [][]int{}
    }
    ans := make([][]int, 0)
    sort.Ints(nums)
    uStart := -1
    for i, num := range nums {
        if num >= 0 {
            uStart = i
            if num == 0 && i+2 < len(nums) && nums[i+2] == 0 {
                ans = append(ans, []int{0, 0, 0})
            }
            break
        }
    }
    // 表示只有负数或正数
    if uStart <= 0 {
        return ans
    }

    for i := 0; i < uStart; i++ {
        // 已经判断过了
        if i-1 >= 0 && nums[i] == nums[i-1] {
            continue
        }
        for left, right := uStart, len(nums)-1; left < right; {
            res := nums[i] + nums[left] + nums[right]
            if res == 0 {
                ans = append(ans, []int{nums[i], nums[left], nums[right]})
                left++
                right--
            } else if res > 0 {
                right--
            } else if res < 0 {
                left++
            }

            if left-1 >= uStart && nums[left] == nums[left-1] && left < right {
                if nums[i]+nums[left]+nums[left-1] == 0 && nums[left] != nums[right] {
                    ans = append(ans, []int{nums[i], nums[left], nums[left-1]})
                }
                for left < right && nums[left] == nums[left-1] {
                    left++
                }
            }

            if right+1 < len(nums) && nums[i] == nums[right+1] && left < right {
                if nums[i]+nums[right]+nums[right+1] == 0 && nums[left] != nums[right] {
                    ans = append(ans, []int{nums[i], nums[right], nums[right+1]})
                }
                for left < right && nums[right] == nums[right+1] {
                    right--
                }
            }
        }
    }

    for i := len(nums) - 1; i >= uStart; i-- {
        if i+1 < len(nums) && nums[i] == nums[i+1] {
            continue
        }
        for left, right := 0, uStart-1; left < right; {
            res := nums[i] + nums[left] + nums[right]
            if res == 0 {
                ans = append(ans, []int{nums[left], nums[right], nums[i]})
                left++
                right--
            } else if res > 0 {
                right--
            } else if res < 0 {
                left++
            }

            if left-1 >= 0 && nums[left] == nums[left-1] && left < right {
                if nums[i]+nums[left]+nums[left-1] == 0 && nums[left] != nums[right] {
                    ans = append(ans, []int{nums[i], nums[left], nums[left-1]})
                }
                for left < right && nums[left] == nums[left-1] {
                    left++
                }
            }

            if right+1 < uStart && nums[right] == nums[right+1] && left < right {
                if nums[i]+nums[right]+nums[right+1] == 0 && nums[left] != nums[right] {
                    ans = append(ans, []int{nums[i], nums[right], nums[right+1]})
                }
                for left < right && nums[right] == nums[right+1] {
                    right--
                }
            }
        }
    }
    return ans
}