0015-3sum.py

"""
15. 3Sum

Submitted: March 2, 2026

Runtime: 1712 ms (beats 7.32%)
Memory: 28.08 MB (beats 5.36%)
"""

from collections import Counter

class Solution:
    def twoSum(self, used: int) -> list[list[int]]:
        if used == 0:
            return [[0, 0]] if self.nums_count[0] >= 3 else []
        elif used < 0:
            # get largest two
            a, b = self.nums_2largest
            # could never add up
            if a + b + used < 0:
                # not possible
                return []
        elif used > 0:
            # smallest two
            a, b = self.nums_2smallest
            # could never add up
            if a + b + used > 0:
                # not possible
                return []

        target = -used
        nums = (
            num for num in self.nums if num != used or self.nums_count[num] > 1
        )
        res = []
        for num in nums:
            diff = target - num
            # to prevent going through the array with two pairs
            if diff > num: continue
            # sanity check here to make sure we don't use numbers too many times
            # used == num == diff != 0 since optimization at top of program
            if used == num == diff:
                continue
            elif used == diff:
                if self.nums_count[used] < 2: continue
            elif used == num or num == diff:
                if self.nums_count[num] < 2: continue

            if num == diff or diff in self.nums_count:
                res.append((num, diff))
        return res

    def threeSum(self, nums: list[int]) -> list[list[int]]:
        if len(set(nums)) == 1:
            return [[0, 0, 0]] if nums[0] == 0 else []

        self.nums = set(nums)
        nums_sorted = sorted(nums)
        self.nums_2largest = nums_sorted[-2:]
        self.nums_2smallest = nums_sorted[:2]
        del nums_sorted
        self.nums_count = Counter(nums)
        res = []
        for num in self.nums:
            twoposs = self.twoSum(num)
            if not twoposs:
                continue
            res.extend((num, *t) for t in twoposs)
        res = sorted(tuple(set(tuple(sorted(t)) for t in res)))
        return res