3SUM by sorting and two-pointer scan in Python

Sort, fix one element, and run a two-pointer scan on the rest, giving O(n^2) time.

File: three_sum_two_pointer_scan.py, 34 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

three_sum_two_pointer_scan.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/three-sum-two-pointer-scan/. The authors are listed in prov/V2066X.jsonl.


def three_sum_two_pointer_scan(values: list[int], t: int = 0) -> list[tuple[int, int, int]]:
    """Return every distinct value triple (x, y, z), x <= y <= z, that is drawn
    from three different positions of values and has x + y + z == t.

    The triples come out in lexicographic order. The input list is not changed.
    Cost: O(n log n) for the sort plus O(n^2) for the scans; O(n) extra space
    for the sorted copy (O(1) more if the caller may sort in place).
    """
    a = sorted(values)  # @line L1
    n = len(a)
    out: list[tuple[int, int, int]] = []
    for i in range(n - 2):  # @line L2
        if i > 0 and a[i] == a[i - 1]:  # @line L3
            continue
        lo, hi = i + 1, n - 1  # @line L4
        while lo < hi:  # @line L5
            s = a[i] + a[lo] + a[hi]  # @line L6
            if s < t:  # @line L7
                lo += 1
            elif s > t:  # @line L8
                hi -= 1
            else:
                out.append((a[i], a[lo], a[hi]))  # @line L9
                lo += 1  # @line L10
                while lo < hi and a[lo] == a[lo - 1]:
                    lo += 1
                hi -= 1  # @line L11
                while lo < hi and a[hi] == a[hi + 1]:
                    hi -= 1
    return out  # @line L12

test_three_sum_two_pointer_scan.py 46 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/three-sum-two-pointer-scan/. The authors are listed in prov/V2066X.jsonl.

import itertools
import random

from three_sum_two_pointer_scan import three_sum_two_pointer_scan


def brute(values: list[int], t: int) -> list[tuple[int, int, int]]:
    found = set()
    for i, j, k in itertools.combinations(range(len(values)), 3):
        if values[i] + values[j] + values[k] == t:
            found.add(tuple(sorted((values[i], values[j], values[k]))))
    return sorted(found)


def test_example() -> None:
    assert three_sum_two_pointer_scan([-1, 0, 1, 2, -1, -4]) == [(-1, -1, 2), (-1, 0, 1)]


def test_edges() -> None:
    assert three_sum_two_pointer_scan([]) == []
    assert three_sum_two_pointer_scan([0, 0]) == []
    assert three_sum_two_pointer_scan([0, 0, 0]) == [(0, 0, 0)]
    assert three_sum_two_pointer_scan([0, 0, 0, 0, 0]) == [(0, 0, 0)]
    assert three_sum_two_pointer_scan([1, 1, 1], 3) == [(1, 1, 1)]
    assert three_sum_two_pointer_scan([1, 2, 4]) == []
    assert three_sum_two_pointer_scan([-2, 0, 0, 2, 2]) == [(-2, 0, 2)]
    big = 2**61
    assert three_sum_two_pointer_scan([-2 * big, big, big]) == [(-2 * big, big, big)]


def test_input_not_changed() -> None:
    a = [3, -1, -2, 0]
    three_sum_two_pointer_scan(a)
    assert a == [3, -1, -2, 0]


def test_against_brute_force() -> None:
    rng = random.Random(7031)
    for _ in range(1500):
        n = rng.randrange(0, 11)
        a = [rng.randrange(-6, 7) for _ in range(n)]
        t = rng.randrange(-5, 6)
        assert three_sum_two_pointer_scan(a, t) == brute(a, t)

What the next level needs

the tests did not run in this build

Read the full page about 3SUM by sorting and two-pointer scan