Exercism pythagorean-triplets: generating canonical-data.json

create_canonical_data.json

#!/usr/bin/env python3.6
import json


def tripletsWithSum(n):
    triplets = set()
    b_limit = int(n / 2) + 1
    for a in range(1, n - 2):
        a2 = a * a
        n_a = n - a
        for b in range(a, b_limit):
            b2 = b * b
            c = n_a - b
            if c > a2 + b2:
                continue
            if c * c == b2 + a2:
                triplets.add((a, b, c))
    return triplets


def nearby(
    n, can_go_down=True, can_go_up=True, increment=1, lower_bound=0,
    upper_bound=0xffffffff
):
    yield n
    x1 = x2 = n
    while can_go_down or can_go_up:
        if can_go_up and x1 < upper_bound:
            x1 += increment
            yield x1
        if can_go_down and x2 > lower_bound:
            x2 -= increment
            yield x2


def create_case_near(
    n, max_triplets=0xffffffff, min_triplets=1,
    can_go_down=True, can_go_up=True,
    fmt='triplets whose sum is {n}'
):
    print(f'Creating case near {n}...', end='', flush=True)
    case = dict(
        description=fmt.format(n=n),
        property='tripletsWithSum',
        input=dict(n=n),
        expected=[]
    )
    triplets = case['expected']
    nearby_nums = nearby(n, can_go_down, can_go_up)
    while True:
        n = next(nearby_nums)
        print(f' {n}', end='', flush=True)
        triplets.extend(tripletsWithSum(n))
        if min_triplets <= len(triplets) and len(triplets) <= max_triplets:
            case['input']['n'] = n
            case['description'] = fmt.format(
                n=n,
                min_triplets=min_triplets,
                max_triplets=max_triplets,
            )
            break
        triplets.clear()
    print('', flush=True)
    print(json.dumps(case, indent=2), flush=True)
    return case


if __name__ == '__main__':
    data = dict(
        exercise='pythagorean-triplet',
        version='1.0.0',
        cases=[
            create_case_near(n, **kwargs)
            for n, kwargs in [
                (10, dict(max_triplets=1)),
                (100, dict(can_go_down=False)),
                (1000, dict()),
                (1000, dict(
                    fmt='no matching triplets for {n}',
                    max_triplets=0,
                    min_triplets=0
                )),
                (90, dict(
                    fmt='returns all matching triplets',
                    min_triplets=2
                )),
                (800, dict(fmt='several matching triplets', min_triplets=7)),
                (30000, dict(fmt='triplets for large number')),
            ]
        ]
    )
    with open('canonical-data.json', 'w') as f:
        f.write(json.dumps(data, indent=4))
    print('Done')