Polynomial parsing and division (by Ruffini's method)

find_roots.py

import poly
import operator
import fractions
import functools
import itertools


def is_root(p, k):
    '''
    True if k is a root of the polynomial p, False otherwise.
    '''

    return not poly.pdivmod(p, k)[1]


def divisors(k):
    '''
    Find all divisors of the number k:

        >>> divisors(1)
        [1]
        >>> divisors(4)
        [1, 2, 4]
        >>> divisors(6)
        [1, 2, 3, 6]
        >>> divisors(12)
        [1, 2, 3, 4, 6, 12]
    '''

    if k == 1:
        return [1]
    pairs = [[i, k/i] for i in xrange(1, int(k ** .5) + 1) if not k % i]
    return functools.reduce(operator.concat, pairs)


if __name__ == '__main__':
    roots = []
    polynomial = poly.parse(raw_input('Polynomial: '))
    # Find divisors of the constant term
    ct_divisors = divisors(abs(polynomial[-1][0]))
    # Find divisors of the highest power coefficient
    hp_divisors = divisors(abs(polynomial[0][0]))
    # By the rational root theorem, check all the possibilities
    for p, q in itertools.product(ct_divisors, hp_divisors):
        for sign in (1, -1):
            potential_root = fractions.Fraction(sign * p, q)
            if is_root(polynomial, potential_root):
                roots.append(potential_root)
    # Check x = 0
    if is_root(polynomial, 0):
        roots.append(0)
    # Output results
    if roots:
        n = len(roots)
        print '%d root%s found:' % (n, ['', 's'][n > 1])
        print ', '.join(map(str, roots))
    else:
        print 'No roots found.'

poly.py

import re

__all__ = ['parse', 'divide']


def _dict2poly(func):
    '''
    Decorator function that transform a dictionary in the form:

        {power: coefficient, ...}

    to a list of tuples representing a polynomial:

        [(coefficient, power), (coefficient, power), ...]
    '''

    def wrapper(*args, **kwargs):
        d = func(*args, **kwargs)
        return sorted(zip(d.values(), d.keys()), key=lambda i: -i[1])
    return wrapper


def _make_poly(list):
    '''
    Build a polynomial from a list of coefficients:

        >>> _make_poly([1, -3, 0, 1]) # x^3 - 3x^2 + 1
        [(1, 3), (-3, 2), (1, 0)]
        >>> _make_poly([1, 0, 0, 0, 1]) # x^4 + 1
        [(1, 4), (1, 0)]
        >>> _make_poly([1, 2, 3, 4]) # x^3 + 2x^2 + 3x + 4
        [(1, 3), (2, 2), (3, 1), (4, 0)]

    The coefficients must be sorted in descending order
    with respect to power.
    '''

    return [(c, p) for c, p in zip(list, xrange(len(list) - 1, -1, -1)) if c]


@_dict2poly
def _fill_coeff(poly):
    '''
    Return a polynomial equal to the given one but with all the powers:

        >>> _fill_coeff([(1, 3), (-1, 0)]) # x^3 - 1
        [(1, 3), (0, 2), (0, 1), (-1, 0)] # x^3 + 0x^2 + 0x - 1
        >>> _fill_coeff([(1, 2)]) # x^2
        [(1, 2), (0, 1), (0, 0)] # x^2 + 0x + 0
    '''

    zero = dict((power, 0) for power in xrange(poly[0][1]))
    zero.update(zip(*zip(*poly)[::-1]))  # update with reversed poly
    return zero


@_dict2poly
def parse(string):
    '''
    Parse a polynomial. You can use ** or ^ to indicate exponentiation.
    Valid polynomials include:

        3x - 2
        x + 1
        4x**2 + x - 1
        -2x^3 + x**2 - x + 1

    Returns a list of tuples (coefficient, power) sorted in descending order
    with respect to power.
    '''

    x_re = re.compile(r'([-+]?\d*)(x?)\^?(\d*)')
    poly = {}
    signs = {'+': 1, '-': -1}
    for c, x, p in x_re.findall(string.replace(' ', '').replace('**', '^')):
        if not any((c, x, p)):
            continue
        coeff = signs[c] if c in ('+', '-') else int(c or 1)
        power = int(p or 1) if x else 0
        # multiple monomials with the same degree
        if power in poly:
            poly[power] += coeff
            continue
        poly[power] = coeff
    return poly


def pdivmod(poly, k):
    '''
    Perform division between poly and (x - k).
    Return a tuple (quotient, remainder).
    '''

    coefficients = zip(*_fill_coeff(poly))[0]
    last = coefficients[0]
    result = [last]
    for c in coefficients[1:]:
        last = c + k * last
        result.append(last)
    quotient = result[:-1]
    remainder = [result[-1]]
    return tuple(map(_make_poly, [quotient, remainder]))

test_poly.py

import unittest
import fractions
from poly import _fill_coeff, _make_poly, parse, pdivmod
from find_roots import is_root, divisors


def _divide_poly(list):
    return pdivmod(*list)


def _test_expected(func):
    def wrapper(meth):
        def inner(self):
            for value, expected in meth(self).iteritems():
                self.assertEqual(func(value), expected)
        return inner
    return wrapper


class TestPoly(unittest.TestCase):

    @_test_expected(_fill_coeff)
    def test_fill_coeff(self):
        return {
                ((3, 2), (-3, 1), (4, 0)): [(3, 2), (-3, 1), (4, 0)],
                ((-3, 3), (-1, 0)): [(-3, 3), (0, 2), (0, 1), (-1, 0)],
                ((0, 2), (1, 1)): [(0, 2), (1, 1), (0, 0)],
                ((3, 2), (-1, 0)): [(3, 2), (0, 1), (-1, 0)],
        }

    @_test_expected(parse)
    def test_parse(self):
        return {
                '3x - 2': [(3, 1), (-2, 0)],
                'x + 1': [(1, 1), (1, 0)],
                '4x**2 + x - 1': [(4, 2), (1, 1), (-1, 0)],
                '-2x^3 + x**2 -x + 1': [(-2, 3), (1, 2), (-1, 1), (1, 0)],
                '- x ^ 3 + 2': [(-1, 3), (2, 0)],
                '4 x': [(4, 1)],
                '- 5 x ^ 3 + 1 - 4': [(-5, 3), (-3, 0)],
                '-x - x^2': [(-1, 2), (-1, 1)],
                'x + x - 3x': [(-1, 1)],
        }

    @_test_expected(_make_poly)
    def test_make_poly(self):
        return {
                (1, -3, 2, 4): [(1, 3), (-3, 2), (2, 1), (4, 0)],
                (-2, 1, 2): [(-2, 2), (1, 1), (2, 0)],
                (1, 0, 1): [(1, 2), (1, 0)],
                (-1, 0, 0, 3, -3): [(-1, 4), (3, 1), (-3, 0)],
                (0, 0, 3): [(3, 0)],
        }

    @_test_expected(_divide_poly)
    def test_divide(self):
        return {
                (((1, 2), (2, 1), (-8, 0)), 2): ([(1, 1), (4, 0)], []),
                (((1, 2), (2, 1), (-8, 0)), -4): ([(1, 1), (-2, 0)], []),
                (((1, 3), (-3, 2), (3, 1), (-1, 0)), 1): \
                        ([(1, 2), (-2, 1), (1, 0)], []),
                (((1, 2), (2, 1), (1, 0)), -2): \
                        ([(1, 1)], [(1, 0)]),
                (((1, 3), (-1, 2), (1, 1)), -4): \
                        ([(1, 2), (-5, 1), (21, 0)], [(-84, 0)]),
        }


class TestFindRoots(unittest.TestCase):

    def test_is_root(self):
        self.assertTrue(is_root([(2, 2), (-5, 1), (3, 0)], 1))
        self.assertTrue(is_root([(2, 2), (-5, 1), (3, 0)],
            fractions.Fraction(3, 2)))
        self.assertFalse(is_root([(2, 2), (-5, 1), (3, 0)], -1))
        self.assertTrue(is_root([(1, 1), (-1, 0)], 1))
        self.assertTrue(is_root([(1, 1)], 0))

    def test_divisors(self):
        self.assertItemsEqual(divisors(12), [1, 2, 3, 4, 6, 12])
        self.assertItemsEqual(divisors(24), [1, 2, 3, 4, 6, 8, 12, 24])
        self.assertItemsEqual(divisors(1), [1])
        self.assertItemsEqual(divisors(11), [1, 11])


if __name__ == '__main__':
    unittest.main()