Leetcode 3sum problem solution

Question

The problem is:

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] where i, j and k are distinct and nums[i] + nums[j] + nums[k] == 0.

The solution set must not contain duplicate triplets and the order of the output and the order of the triplets does not matter.

This is my code:

def threeSum(nums):
    triples = set() # this is where the triples will be stored
    num_map = {}
    n = len(nums)

    for i in range(n): # creating a hash map of the array for quick lookups 
        num_map[nums[i]] = i 

    for i in range(n-1):
        for j in range(i+1,n):
            comp = - nums[i] - nums[j] # the complement
            if (comp in num_map and num_map[comp] != i and num_map[comp] != j):
                triples.add(tuple(sorted([nums[i],nums[j],comp])))
    return triples

My questions are:

1. I believe the time complexity of my solution is O(n^2), and it is not known if there is an algorithm which can perform better than O(n^2). Am I correct?
2. In any event, it would seem that my code is not efficient. Which part of the code is slowing it down?

I suspected the use of sorted() might be adding to it, but given that it is only sorting an array of 3 elements (and is only triggered if a solution is found), I thought it might not be having much impact. The reason I sort the array is to avoid duplicates.

Answer 1

It would be nice to see some tests of the function. That would help reviewers see what's known to work, and identify untested scenarios.

---

This transformation is strongly tied to the specific requirements of the question:

for i in range(n): # creating a hash map of the array for quick lookups 
    num_map[nums[i]] = i

We're storing only a single index in this reverse lookup table, and overwriting any previous value. The wording of the question permits this because we are required to return a set of the values. But this approach makes the code less reusable for similar problems (such as those where we must return the indices rather than values).

---

for i in range(n) is unidiomatic. I'd write:

    for i,n in enumerate(nums):
        num_map[n] = i
        print(num_map)

Or more likely:

    num_map = {n:i for i,n in enumerate(nums)}

---

Since we don't need to know the actual indices used (only that they be distinct), we can transform the input by sorting it input first, thus allowing us to know that nums[i] ≤ nums[j] ≤ nums[k] when i ≤ j ≤ k. Then we don't need sorted when adding, and our condition can be simpler:

    nums = sorted(nums)
    num_map = dict((n,i) for i,n in enumerate(nums))

    for i,nᵢ in enumerate(nums[:-1]):
        for j,nⱼ in enumerate(nums[i+1:], i+1):
            comp = -nᵢ -nⱼ # the complement
            if comp in num_map and j < num_map[comp]:
                triples.add((nᵢ, nⱼ, comp))

We can then reduce work by observing that when we reach nⱼ that would make k too small, then no higher j will succeed, so we can break out of the j loop early:

    for i,nᵢ in enumerate(nums[:-1]):
        for j,nⱼ in enumerate(nums[i+1:], i+1):
            comp = -nᵢ -nⱼ # the complement
            if nⱼ > comp:
                # we won't find any more matches in this j or higher
                break
            try:
                if j < num_map[comp]:
                    triples.add((nᵢ, nⱼ, comp))
            except KeyError:
                pass

I also eliminated the duplication of comp in num_map and num_map[comp] here - remember that Python exceptions are cheap.

We could probably go further, and remove values from comp as they become unreachable. Or maintain an index into it, searching backwards as j increases, rather than doing the in test each iteration. I'll leave that as an exercise for the motivated reader.

Answer 2

The excellent answer by user @Toby Speight addresses the main goal of performance.

I'm not that smart, but I'll offer some feedback on more mundane style issues with your code.

Overview

The code layout is nice and compact.

Documentation

It would be nice to add a docstring to the function to describe ts purpose, what the input is and what it returns. The problem statement in your question seems sufficient.

Comments

The following comment is not very useful because it does not tell the user anything new; it is obvious you are storing something into a variable.

# this is where the triples will be stored

A "hash map" in Python parlance is a "dictionary".

Naming

The variable triples could be named triplets to be consistent with the documentation.

Lint check

pylint identified trailing whitespace on a couple lines. Since it is not needed and can be quite annoying, just delete it.

While we're on the topic, I think adding a single space after each comma makes the code easier to read.

---

Here is new code with the suggestions above:

def threeSum(nums):
    '''
    Given an integer array nums, return all the triplets
    [ nums[i], nums[j], nums[k] ] where i, j and k are distinct
    and nums[i] + nums[j] + nums[k] == 0.
    '''
    triplets = set()
    num_map = {}
    n = len(nums)

    for i in range(n): # creating a dictionary of the array for quick lookups
        num_map[nums[i]] = i

    for i in range(n-1):
        for j in range(i+1, n):
            comp = - nums[i] - nums[j] # the complement
            if (comp in num_map and num_map[comp] != i and num_map[comp] != j):
                triplets.add(tuple(sorted([nums[i], nums[j], comp])))
    return triplets