15. 3Sum

• categories: Code, Interview Question, leetcode, Medium
• source: https://leetcode.com/problems/3sum
• topics: Hash Table

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.

Notice that the solution set must not contain duplicate triplets.

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        ans = set()
        index = {}
        nums.sort()
 
        # Save max index of pivot = candidates for nums[k]
        for i in range(len(nums)-1, -1, -1):
            if not -nums[i] in index:
                index[-nums[i]] = i
 
        prev = float('-inf')
 
        # Loop through sums respecting that i < j < k
        for i in range(len(nums)):
            # Don't use same element twice and filter positiv values
            if nums[i] > 0 or nums[i] == prev:
                continue
            for j in range(i+1, len(nums)):
                two_sum = nums[i] + nums[j]
                if (not two_sum in index or 
                    index[two_sum] <= j or 
                    (nums[i], nums[j]) in ans):
                    continue
                else:
                    ans.add((nums[i], nums[j]))
            prev = nums[i]
        return [[x[0], x[1], -x[0]-x[1]] for x in ans]