Ruth-Aaron Pairs - Dailyprogrammers [2015-10-05] Challenge #235 [Easy]

dlyprgrmr235e.py

#!/usr/bin/env python3

import unittest


def build_sieve(n):
    """Create the extended and optimized Sieve of Eratosthenes.

    The size of the sieve is `N`
    `sieve[i]`` is `0` if `i` is prime
    else, `sieve[i]`` is the minor prime factor of `i`.
    """
    N = n + 1
    sieve = [0] * N

    for i in range(4, N, 2):
        sieve[i] = 2

    for i in range(3, N, 2):
        if sieve[i] == 0:
            for j in range(i * i, N, 2 * i):
                sieve[j] = i
    return sieve

class Factorize:
    """Factorization function based on an Sieve of Eratosthenes."""

    def __init__(self, N=10000):
        """Build the size `N` sieve to be used by the factorization."""
        self.sieve = build_sieve(N)

    def __call__(self, n, distinct=True):
        """Implementation of the factorization function.

        returns a collection of the prime factors of n.
        If distinct is set to True, no repeated factors will be included
        """
        factors = set() if distinct else []
        insert = set.add if distinct else list.append

        while self.sieve[n] != 0:
            insert(factors, self.sieve[n])
            n //= self.sieve[n]

        insert(factors, n)
        return factors


class RuthAaron:
    """Function that determine if two numbers are Ruth-Aaron.

    Two numbers are Ruth-Aaron if they are a consecutive pair such as their sum
    of distinct prime factors are equal.
    """

    def __init__(self, factorize):
        """Determine with what facotrization function work."""
        self.factorize = factorize

    def __call__(self, n1, n2):
        """Implementation to the Ruth-Aaron Test."""
        return abs(n1 - n2) == 1 and \
            sum(self.factorize(n1)) == sum(self.factorize(n2))


def main():
    messages = {
        True: 'VALID',
        False: 'NOT VALID'
    }
    pairs = []

    import ast
    cases = int(input())
    for case in range(cases):
        pairs.append(ast.literal_eval(input()))

    factorize = Factorize(max(_[1] for _ in pairs))
    ruth_aaron = RuthAaron(factorize)

    for pair in pairs:
        print(pair, messages[ruth_aaron(*pair)])


class TestSieve(unittest.TestCase):

    def setUp(self):
        self.known_sieve = [0, 0, 0, 0, 2, 0, 2, 0, 2,
                            3, 2, 0, 2, 0, 2, 3, 2, 0, 2, 0, 2]

    def test_known_sieve(self):
        sieve = build_sieve(20)
        self.assertEqual(self.known_sieve, sieve)


class TestFactorize(unittest.TestCase):

    def setUp(self):
        self.known_factors = {
            2: set([2]),
            3: set([3]),
            4: set([2]),
            5: set([5]),
            6: set([2, 3]),
            77: set([7, 11]),
            78: set([2, 3, 13]),
            714: set([2, 3, 7, 17]),
            715: set([5, 11, 13]),
        }
        self.factorize = Factorize(max(n for n in self.known_factors.keys()))

    def test_known_factorization_distinct(self):
        for n, factors in self.known_factors.items():
            self.assertEqual(self.factorize(n, distinct=True), factors)


class TestRuthAaron(unittest.TestCase):

    def setUp(self):
        self.known_ruth_aaron = [
            (77, 78),
            (714, 715),
            (5, 6),
            (2107, 2108),
            (492, 493)
        ]

        self.know_not_ruth_aaron = [
            (20, 21),
            (128, 129),
            (6, 7)
        ]

        self.factorize = Factorize(max(
            n for n in
            [_[1] for _ in self.known_ruth_aaron] +
            [_[1] for _ in self.know_not_ruth_aaron]
        ))

        self.ruth_aaron = RuthAaron(self.factorize)

    def test_known_ruth_aaron(self):
        for pair in self.known_ruth_aaron:
            self.assertTrue(self.ruth_aaron(*pair))


if __name__ == '__main__':
    main()