"""
15. 3Sum
Medium

5344

648

Add to List

Share
Given an array nums of n integers, are there elements a, b, c in nums such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note:

The solution set must not contain duplicate triplets.

Example:

Given array nums = [-1, 0, 1, 2, -1, -4],

A solution set is:
[
  [-1, 0, 1],
  [-1, -1, 2]
]
"""


class Solution:
    def threeSum(self, nums):
        pairs = []
        nums.sort()
        for i in range(len(nums) - 2):
            if i > 0 and nums[i] == nums[i - 1]:
                continue
            if nums[i] > 0:
                break
            l, r = i + 1, len(nums) - 1
            while l < r:
                threeSum = nums[i] + nums[l] + nums[r]
                # 요 매번 생성이 속도 저하를 불러옴 아래와 차이가 많이 남
                pair = [nums[i], nums[l], nums[r]]
                if threeSum > 0:
                    r -= 1
                elif threeSum < 0:
                    l += 1
                else:
                    pairs.append(pair)
                    while l < r and nums[l] == nums[l + 1]:
                        l += 1
                    while l < r and nums[r] == nums[r - 1]:
                        r -= 1
                    l += 1
                    r -= 1
        return pairs

    #faster
    def threeSum(self, nums):
        res = []
        nums.sort()
        for i, v in enumerate(nums[:-2]):
            if v > 0: break
            if i > 0 and v == nums[i - 1]: continue
            l = i + 1
            r = len(nums) - 1
            while l < r:
                total = nums[l] + v + nums[r]

                if total < 0:
                    l += 1
                elif total > 0:
                    r -= 1
                else:
                    res.append([nums[l], v, nums[r]])
                    while l < r and nums[l] == nums[l + 1]:
                        l += 1
                    while l < r and nums[r] == nums[r - 1]:
                        r -= 1
                    l += 1
                    r -= 1
        return res


S = Solution()
nums = [-1, 0, 1, 2, -1, -4]
test = S.threeSum(nums)
print(test)
