find_k_sum.py

#!/usr/bin/env python3

from timeit import timeit

def naive(values, k):
    for i, v1 in enumerate(values):
        for v2 in values[i + 1:]:
            if v1 + v2 == k:
                return True
    return False

def better(values, k):
    return any(
        (k - m in values[i + 1:])
        for i, m in enumerate(values)
    )

def best(values, k):
    seen = set()
    return any(
        k - v in seen or seen.add(v)
        for v in values
    )

def tester(fn):
    values = list(range(10))
    assert not fn(values, 0), "0-test"
    assert not fn(values, 2 * values[-1]), "Last test"
    for i in range(1, 2 * values[-1]): 
        assert fn(values, i), f"test for {i}"

def timer():
    number = 1000
    values = list(range(number))
    glob = {
        "naive": naive,
        "better": better,
        "best": best,
        "values": values,
        "k": 2 * number,
    }
    for fn in (naive, better, best):
        print(fn.__name__)
        print(timeit(stmt=f"{fn.__name__}(values, k)", globals=glob, number=100))

def test():
    print("naive")
    tester(naive)
    print("better")
    tester(better)
    print("best")
    tester(best)

if __name__ == '__main__':
    # test()
    timer()