p-adic-numbers/poly.py

# Polynomial class on Z or Z_p

from collections import defaultdict

class Poly:
    'Polynomial class.'
    def __init__(self, coeffs = None):
        self.coeffs = defaultdict(int, isinstance(coeffs, int) and {0:coeffs} or coeffs or {})
        self.deg = int(len(self.coeffs) and max(self.coeffs.keys()))
    
    def __call__(self, val):
        'Evaluate polynomial for a given value.'
        res = 0
        for i in xrange(self.deg, -1, -1):
            res = res * val + self.coeffs[i]
        return res
    
    def __str__(self):
        def term(coeff, expt):
            if coeff == 1 and expt == 0:
                return '1'
            return ' * '.join(([] if coeff == 1 else [str(coeff)]) + \
                              ([] if expt == 0 else ['X'] if expt == 1 else ['X ** %d' % expt]))
            
        return ' + '.join(term(self.coeffs[i], i) for i in self.coeffs if self.coeffs[i] != 0)
    def __repr__(self):
        return str(self)
    
    # arithmetic
    def __neg__(self):
        return Poly({(i, -self.coeffs[i]) for i in self.coeffs})
    def __add__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        res = Poly()
        res.deg = max(self.deg, other.deg)
        for i in xrange(res.deg+1):
            res.coeffs[i] = self.coeffs[i] + other.coeffs[i]
        return res
    def __radd__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return other.__add__(self)
    def __sub__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return self.__add__(other.__neg__())
    def __rsub__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return other.__add__(self.__neg__())
    
    def __mul__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        res = Poly()
        res.deg = self.deg + other.deg # consider case where other is 0
        for i in xrange(res.deg+1):
            for j in xrange(i+1):
                res.coeffs[i] += self.coeffs[j] * other.coeffs[i - j]
        return res
    def __rmul__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return other.__mul__(self)
        
    def __pow__(self, other):
        if not isinstance(other, int) or other < 0:
            raise ValueError("Exponent %d is not a natural number" % other)
        res = Poly(1)
        while other:
            res *= self
            other -= 1
        return res

X = Poly({1:1})

def derivative(p):
    'Return derivative of polynomial.'
    return Poly({(i - 1, i * p.coeffs[i]) for i in p.coeffs if i != 0})