3Sum

From LeetCode

problem description / no solution provided

Solution in Python3

from collections import Counter
from bisect import bisect, bisect_left

class Solution:
    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        triplets = []
        if len(nums) < 3:
            return triplets
        num_freq = Counter(nums)
        nums = sorted(num_freq)  # Sorted unique numbers
        max_num = nums[-1]
        for i, v in enumerate(nums):
            if num_freq[v] >= 2:
                complement =  -2 * v
                if complement in num_freq:
                    if complement != v or num_freq[v] >= 3:
                        triplets.append([v] * 2 + [complement])

            # When all 3 numbers are different.
            if v < 0:  # Only when v is the smallest
                two_sum = -v

                # Lower/upper bound of the smaller of remaining
                # two.
                lb = bisect_left(nums, two_sum - max_num, i + 1)
                ub = bisect(nums, two_sum // 2, lb)
                       
                for u in nums[lb : ub]:
                    complement = two_sum - u
                    if complement in num_freq and u != complement:
                        triplets.append([v, u, complement])
        return triplets

I am lucky to get

Runtime: 400 ms, faster than 100% of Python3 online submissions for 3Sum.

Time Complexity

\(O(n^2)\), where \(n\) is the number of numbers.

Variants

Keep balanced

        nums = sorted(num_freq)  # Sorted unique numbers
+
+       # Get rid of numbers that are too large/small such that
+       # no other number able to complete.
+       lbound = bisect_left(nums, -2 * nums[-1])
+       if lbound != 0:  # Some numbers are too small  
+           nums = nums[lbound:]
+       else:
+           nums = nums[:bisect(nums, -2 * nums[0])]
+       if len(nums) < 1:
+           return triplets
+       
        max_num = nums[-1]

I am lucky to get

Runtime: 316 ms, faster than 100% of Python3 online submissions for 3Sum.

Two pointers

            # When all 3 numbers are different
            if v < 0:  # only when v is smallest
                two_sum = -v

-               # Lower/upper bound of the smaller of remaining
-               # two.
-               lb = bisect_left(nums, two_sum - max_num, i + 1)
-               ub = bisect(nums, two_sum // 2, lb)
+               # Lower bound for the smaller of remaining two.
+               lb = bisect_left(nums, two_sum - max_num, i + 1)
+               lb = min(lb, len(nums) - 1)
+
+               # Upper bound of the greater of remaining two.
+               ub = bisect(nums, two_sum - nums[lb], lb + 1)
+               ub = min(ub, len(nums) - 1)
                      
-               for u in nums[lb : ub]:
-                   complement = two_sum - u
-                   if complement in num_freq and u != complement:
-                       triplets.append([v, u, complement])
+               while lb < ub:
+                   if nums[lb] + nums[ub] == two_sum:
+                       triplets.append([v, nums[lb], nums[ub]])
+                       lb += 1
+                       ub -= 1
+                   elif nums[lb] + nums[ub] > two_sum:
+                       ub -= 1
+                   else:
+                       lb += 1
        return triplets

This solution is not bad, it has:

Runtime: 808 ms, faster than 96.57% of Python3 online submissions for 3Sum.

In case we can’t get binary search effortlessly, let lb = i + 1 and ub = len(nums) - 1. This sample has

Runtime: 492 ms, faster than 99.65% of Python3 online submissions for 3Sum.

Not bad at all.

Remove duplicated triplets

This problem is difficult only because we need many conditions to keep all triplets are unique (otherwise solution will be simple). If we handle duplicates seperately, it could be easier to write and read. However, it may not worth the effort to compare triplets given the data structure we have by hand. It will also sacrifice runtime.

Sum up to arbitrary target

It is not hard to make our program more general, for example

-           if v < 0:  # Only when v is smallest
-               two_sum = -v
+           if v * 3 < target:  # Only when v is smallest
+               two_sum = target - v

This is not the only piece needs modification, but you get the idea. sample1 sample2