boing.py

#!/usr/bin/env python3

import pytest


def boing(mass_factor):
    little_v = 0
    big_v = -1

    bounces = 0

    while not (0 <= little_v < big_v and big_v >= 0):
        (little_v, big_v) = bounce(1, little_v, mass_factor, big_v)
        bounces += 1
        if little_v < 0:
            little_v = -1 * little_v
            bounces += 1
    return bounces


def bounce(m1, v1, m2, v2):
    new_v1 = (((m1 - m2) * v1) / (m1 + m2)) + ((2 * m2 * v2) / (m1 + m2))
    new_v2 = (((2 * m1) * v1) / (m1 + m2)) + (((m2 - m1) * v2) / (m1 + m2))
    return new_v1, new_v2


@pytest.mark.parametrize("m1, v1, m2, v2, expected_v1, expected_v2", [
    (1, -1, 1, 0, 0, -1),
    (1, 0, 1, -1, -1, 0),
    (1, 1, 1, -1, -1, 1)
])
def test_bounce(m1, v1, m2, v2, expected_v1, expected_v2):
    assert (expected_v1, expected_v2) == bounce(m1, v1, m2, v2)


@pytest.mark.parametrize("mass_factor,expected", [
    (1, 3),
    (100, 31),
    (10000, 314),
    (1000000, 3141),
])
def test_boing(mass_factor, expected):
    assert expected == boing(mass_factor)


@pytest.mark.parametrize("num_digits,expected", [
    (1, 3),
    (2, 31),
    (3, 314),
    (4, 3141),
])
def test_pi_digits_compute(num_digits, expected):
    assert expected == boing(100 ** (num_digits - 1))


if __name__ == "__main__":
    import sys
    num_digits = int(sys.argv[1])
    print(boing(100 ** (num_digits - 1)))