Two-Sum with Generator Exceeds Time-CodePudding

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?

CodePudding user response：

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.