Two Sum

Saturday, February 22, 2025

A few days ago, Ryan Petermen wrote a tweet about switching to Python from C++ in a solution for the Two Sum problem from LeetCode:

Given an array of integers nums and an integer target, return indices of the two numbers such that they add up to target.

You may assume that each input would have exactly one solution, and you may not use the same element twice.

You can return the answer in any order.

Spoilers below!

They offered this iterative C++ solution:

std::vector<int> twoSum(std::vector<int>& nums, int target) {
    std::unordered_map<int, int> hashMap;
    std::vector<int> result;

    for (int i = 0; i < nums.size(); ++i) {
        int complement = target - nums[i];
        if (hashMap.find(complement) != hashMap.end()) {
            result.push_back(hashMap[complement]);
            result.push_back(i);
            return result;
        }
        hashMap[nums[i]] = i;
    }
    return result;
}

Followed by this improved Python solution:

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

    for i, num in enumerate(nums):
        complement = target - num
        if complement in hashmap:
            return [hashmap[complement], i]
        hashmap[num] = i

Of course, I was curious what a Factor solution would look like.

Direct translation

We can start by making a mostly direct translation using local variables:

:: two-sum ( nums target -- i j )
    H{ } clone :> hashmap
    nums <enumerated> [
        first2 :> ( i num )
        target num - :> complement
        complement hashmap at [
            i num hashmap set-at f
        ] unless*
    ] map-find ?first ;

And a few test cases to show that it works:

{ 0 1 } [ { 2 7 11 15 } 9 two-sum ] unit-test
{ 1 2 } [ { 3 2 4 } 6 two-sum ] unit-test
{ 0 1 } [ { 3 3 } 6 two-sum ] unit-test

Using map-find-index

Sometimes a higher-level word exists that can simplify logic, for example map-find-index:

:: two-sum ( nums target -- i j )
    H{ } clone :> hashmap
    nums [| num i |
        target num - :> complement
        complement hashmap at [
            i num hashmap set-at f
        ] unless*
    ] map-find-index drop ;

Implicit arguments

In the spirit of concatenative thinking, we could reduce the amount of named variables we have slightly by not naming the complement internal variable:

:: two-sum ( nums target -- i j )
    H{ } clone :> hashmap
    nums [| num i |
        target num - hashmap at [
            i num hashmap set-at f
        ] unless*
    ] map-find-index drop ;

Fried quotations

And that works, but perhaps we could try some alternative syntax features like fried quotations:

: two-sum ( nums target -- i j )
    H{ } clone dup '[
        swap _ over - _ at [ 2nip ] [ _ set-at f ] if*
    ] map-find-index drop ;

Is that better? Perhaps.

How else might you solve this problem?

Can you do it in one-line?