Hash Map O(1) Lookups

Drop the sorted guarantee. Track lengths now arrive in play order, and the task is to return the positions of 2 that add up to a target. The direct answer tries every pair with a nested loop:

def two_sum_pairs(nums, target):
    for i in range(len(nums)):
        for j in range(i + 1, len(nums)):
            if nums[i] + nums[j] == target:
                return [i, j]
    return None

Starting the inner loop at i + 1 tries each pair once, yet the count of pairs grows with the square of the list length: double the list and the worst case roughly quadruples.

Turn the search into a lookup

For any number, its partner is fully decided: it must equal target - num. The question 'which number pairs with this one?' becomes 'have I already seen the value target - num?' A dict answers that with one hash lookup, and the cost of a lookup stays flat whether the dict holds 4 keys or 4 million. The keys are values seen so far, and each maps to the position where it appeared.

def two_sum(nums, target):
    seen = {}

    for i, num in enumerate(nums):
        complement = target - num

        if complement in seen:
            return [seen[complement], i]

        seen[num] = i

    return None

nums = [2, 7, 11, 15]
print(two_sum(nums, 9))

[0, 1]

At position 0 the loop needs 7, which is not in the dict yet, so it stores {2: 0}. At position 1 it needs 2, which is a key, so it returns [0, 1] and never reads 11 or 15. On a list with no answer, each number is looked up once and stored once, so the work is one pass plus one dict entry per item.

Look up first, then store

The order of the 2 dict operations inside the loop decides correctness. Store the current number before checking, and a number can pair with itself:

def two_sum_store_first(nums, target):
    seen = {}
    for i, num in enumerate(nums):
        seen[num] = i
        if target - num in seen:
            return [seen[target - num], i]
    return None

print(two_sum_store_first([5, 3, 8], 10))

Output:

[0, 0]

The list [5, 3, 8] has no pair that sums to 10, yet the function reports positions 0 and 0: the 5 at position 0 found itself in the dict. Checking first means seen holds only earlier positions when the lookup runs. Repeated values still work in that order: on [5, 5] with a target of 10, position 1 needs 5, finds the 5 stored at position 0, and returns [0, 1].

Grouping is the same move

The nested-loop habit appears outside pair problems too. Splitting play events by user with an outer loop over users and an inner filter over all events reads every event once per user. A dict keyed by user does the split in one pass, sending each event straight to its user's list:

from collections import defaultdict

plays = [("ana", "t1"), ("ben", "t4"), ("ana", "t7"), ("cy", "t1"), ("ben", "t2")]
by_user = defaultdict(list)
for user, track in plays:
    by_user[user].append(track)

print(dict(by_user))
print(len(by_user["ana"]))

Output:

{'ana': ['t1', 't7'], 'ben': ['t4', 't2'], 'cy': ['t1']}
2

What does this print?

def two_sum(nums, target):
    seen = {}
    for i, num in enumerate(nums):
        complement = target - num
        if complement in seen:
            return [seen[complement], i]
        seen[num] = i
    return None

print(two_sum([3, 2, 4], 6))

• [0, 0], because 3 is its own partner for a target of 6
• [0, 2], because 3 and 4 sit at the first and last positions
• None, because the value 6 never appears in the list
• [1, 2], because 4 at position 2 finds 2 stored at position 1

When the answer needs only whether a value appeared, the positions stored in the dict go unused, and a set does the same work with keys alone. x in seen is the same flat-cost hash lookup, and seen.add(x) records a value. Finding the first track id played a second time is the two-sum loop with a set: if the current id is already in seen, print it and stop; if not, add it. The order matters for the same reason as before: add first, and every id finds itself on its first visit.

About this lesson

This section is part of the Problem Solving: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.