Two-Sum with Generator Exceeds Time

Question

I'm practicing some problems and started trying different solutions for the two-sum problem.

My first solution is the following:

nums = [2,7,11,15]
target = 9

def twoSum(nums: List[int], target: int) -> List[int]:
    """Given an array and a target, return the indexes of the elements adding to the target"""
    nums_len = len(nums)
    result = []
    for first in range(nums_len):
        for second in range(first + 1, nums_len):
            if len(result) == 2: break
            if nums[first] + nums[second] == target:
                result.extend([first, second])
        else:
            continue
        break
    return result
    
twoSum(nums, target) # returns [0,1] (sum of 2 and 7 = 9)

This solution got a runtime and memory usage of 8385 ms and 15.1 MB, respectively.

When trying to improve the code and readability, I decided to use a generator instead of having those else, continue, break statements at the end of the loops. Here's the code:

def double_loop(arr):
    for i in range(len(arr)):
        for j in range(i + 1, len(arr)):
            yield [i, j]
        
def twoSum(nums: List[int], target: int) -> List[int]:
    for result in double_loop(nums):
        if nums[result[0]] + nums[result[1]] == target:
            break
    return result

nums = [2,7,11,15]
target = 9
twoSum(nums, target) # returns the same as above

The issue is when I submit to Leetcode and it runs different inputs. Especially one where nums = [1, 2, 3, ..., 9999,10000] and target = 19999. I get a Time Limit Exceeded error.

Why is the second solution apparently wasting more resources and timing out? Can someone explain what's wrong with that approach?

Answer 1

Your efficiency question about generators has been answered. But if your goal is to improve readability, look for what you can eliminate from your solution. Take your original solution:

def twoSum(nums: List[int], target: int) -> List[int]:
    """Given an array and a target, return the indexes of the elements adding to the target"""
    nums_len = len(nums)
    result = []
    for first in range(nums_len):
        for second in range(first + 1, nums_len):
            if len(result) == 2: break
            if nums[first] + nums[second] == target:
                result.extend([first, second])
        else:
            continue
        break
    return result

If you look, including result just makes your code more complex for no return. The same thing can be streamlined to:

def twoSum(nums: List[int], target: int) -> List[int]:
    """Given an array and a target, return the indexes of the elements adding to the target"""
    nums_len = len(nums)
    for first in range(nums_len):
        for second in range(first + 1, nums_len):
            if nums[first] + nums[second] == target:
                return [first, second]

This did change behavior slightly. If there is no solution you now return None instead of []. That is arguably more correct. More importantly, the code directly describes what it is doing. This will make life easier for anyone else who is reading it.

And if you want to improve performance, the biggest win is from learning more about data structures rather than refactoring to use more features. The following should scale much better.

def twoSum(nums: List[int], target: int) -> List[int]:
    """Given an array and a target, return the indexes of the elements adding to the target"""
    nums_len = len(nums)
    seen_at = {}
    for second in range(nums_len):
        if target - nums[second] in seen_at:
            return [seen_at[target - nums[second]], second]
        else:
            seen_at[nums[second]] = second

The reason that this improves things is that rather than a second loop to find the other summand, we just do a lookup for whether we've seen it. With n things, this is a O(n) algorithm instead of O(n^2).

On your time limit exceeded case, this runs in under 40 ms on my laptop.

Answer 2

In the second example, you've replaced a relatively straight forward double nested loop (which is cheap, it just has two loop variables and some logic to set and increase them) with a generator that internally has the exact same logic, but adds the overhead of yielding those pairs of values one at a time and then receiving and unpacking them to work with them.

So, the expectation is definitely that it will be slower, you're doing more work after all.

There are some places to optimise your code though:

• Python is often more efficient if you iterate directly over an iterable, like the list in this case (instead of ranging an index and indexing later)
• you could choose to yield your results instead of returning a list, if there's a possibility the user of your function is not interested in all the results, but just the first (few)

So, here's a better version of your first solution:

# you use List, but since 3.9, you could just use list and avoid this import
from typing import List, Generator


nums = [2, 7, 11, 15]
target = 9


def twoSum(nums: List[int], target: int) -> Generator[int, None, None]:
    for i, first in enumerate(nums):
        for j, second in enumerate(nums[i+1:], start=i+1):
            if first + second == target:
                yield (i, j)  # yielding a tuple instead of a list, which suits

# in this case we want all the results, so grabbing them in a list
result = list(twoSum(nums, target))
print(result)

Output:

[(0, 1)]

An example of using the generator for a less manageable result:

big_result = twoSum(list(range(100000)), 237)
print(next(big_result), next(big_result))

And here's a version of your function that fixes two additional points: adding back in the limit to return a needed number of results, and being more permissive about the iterable passed in (so that you don't have to cast to a list to avoid type hints):

from typing import Iterable, Generator


def twoSum(nums: Iterable[int], target: int, max_results: int=-1) -> Generator[int, None, None]:
    for i, first in enumerate(nums):
        for j, second in enumerate(nums[i+1:], start=i+1):
            if not max_results:
                return
            if first + second == target:
                if max_results > 0:
                    max_results -= 1
                yield (i, j)


needed_result = list(twoSum([2, 7, 11, 15], 9, 1))

But also note that you wouldn't normally put in a limit like that, but rather just get the number of results from the generator and then discard it. The generator will get garbage-collected once there's no reference to it anymore.

Answer 3

It's possible to reduce the complexity to O(n log n).

You need first to sort the nums list. Then, for each n in nums, you binary search the list for target-n (binary search costs O(log n) ).

def binary_search(self, a, x, lo=0, hi=None):
  """ ref: https://stackoverflow.com/questions/212358/ """
  from bisect import bisect_left
  if hi is None: 
    hi = len(a)
  pos = bisect_left(a, x, lo, hi)                  # find insertion position
  return pos if pos != hi and a[pos] == x else -1  # don't walk off the end

def twoSum(self, nums: List[int], target: int) -> List[int]:
  nums.sort()
  for i,n in enumerate(nums):
    j = self.binary_search(nums, target-n)
    if j != -1:
      return [i,j]