unreasonably effective

you can be sloppy, as long as you are rigorous

LeetCode 15: 3Sum

link

Time: \mathcal{O}(n^3), space: \mathcal{O}(1).

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        triplets = set()
        for i, first in enumerate(nums[0:n-2]):
            for j, second in enumerate(nums[i+1:n-1], start=i+1):
                for third in nums[j+1:]:
                    if first+second+third == 0:
                        triplets.add( tuple(sorted([first, second, third])) )
        return list(triplets)

Sorting helps in two ways:

  1. We can skip duplicate first and second elements which ensures distinct third.
  2. For a given first element, we can now find the second and third elements in one scan.

Time: O(n^2), space: sorting’s worst-case, \mathcal{O}(n).

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        nums.sort()
        triplets = []
        for i, first in enumerate(nums):
            # Use distinct first element
            if i > 0 and nums[i-1] == nums[i]:
                continue
        
            j, k = i+1, len(nums)-1
            while j < k:
                # Use distinct second element
                if j > i+1 and nums[j-1] == nums[j]:
                    j += 1
                    continue

                triplet_sum = first + nums[j] + nums[k]
                if triplet_sum == 0:
                    triplets.append([ first, nums[j], nums[k] ])
                    j += 1
                    k -= 1
                    continue
                
                # Try increasing the sum
                if triplet_sum < 0:
                    j += 1
                else: # Try decreasing the sum
                    k -= 1

        return triplets

Instead of sorting, we could use one set to de-duplicate first elements and a second set to to find the (second, third) pair in a single scan.

Time: \mathcal{O}(n^2), space: \mathcal{O}(n).

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        triplets = set()
        dup_first = set()
        for i, first in enumerate(nums):
            if first in dup_first:
                continue
            dup_first.add(first)
            target = -first
            seen = set()
            for second in nums[i+1:]:
                if (third := target-second) in seen:
                    triplets.add( tuple( sorted( [first, second, third] ) ) )
                seen.add(second)
        return list(triplets)

tanvirdotzaman

Leave a comment