DALGLI0/poly.py

from fractions import gcd
class Polynomial():
    def __init__(self, lst, mod):
        self.poly = list(map(lambda x: x % mod, lst))    
        self.mod = mod
        self.normalize()
    def normalize(self):
        while self.poly and self.poly[-1] == 0:
            self.poly.pop() 

    #zwraca jednomian stopnia n 
    @staticmethod
    def Monomial(n, c, mod):
        zeros = [0]*n

        zeros.append(c)
        return Polynomial(zeros, mod)

    def __add__(self, p2):
        p1 = self
        len_p1, len_p2= len(p1.poly), len(p2.poly)
        res = [0] * max(len_p1, len_p2)
        if len_p1 > len_p2:
            for _ in range(len_p1-len_p2):
                p2.poly.append(0)
        else:
            for _ in range(len_p2-len_p1):
                p1.poly.append(0)

        for i in range(len(res)):
            res[i] = (p1.poly[i] + p2.poly[i]) % self.mod
        return Polynomial(res, self.mod)

    def __sub__(self, p2):
        p1 = self
        res = []
        len_p2 = len(p2.poly)
        for i in range(len(p1.poly)):
            if i < len_p2:
                res.append(p1.poly[i] - p2.poly[i] % self.mod)
            else:
                res.append(p1.poly[i])
        return Polynomial(res, self.mod)

    def __mul__(self, p2):
        res = [0]*(len(self.poly)+len(p2.poly)-1)
        for i, x1 in enumerate(self.poly):
            for j, x2 in enumerate(p2.poly):
                res[i+j] += x1 * x2 % self.mod
        return Polynomial(res, self.mod)

    def __eq__(self, p2):
        p1 = self
        return p1.poly == p2.poly and p1.mod == p2.mod

    def __pow__(self, n):
        p1 = self
        for i in range(n):
            p1 = p1 * p1
        return p1
    
    def __mod__(self, p2):
        p1 = self
        m = self.mod
        
        if len(p1.poly) < len(p2.poly):
            return p1
        
        if len(p2.poly) == 0:
            raise ZeroDivisionError
            
        divisor_coeff = p2.poly[-1]
        divisor_exp = len(p2.poly) - 1 
        while len(p1.poly) >= len(p2.poly):   
            max_coeff_p1 = p1.poly[-1] #wspolczynnik przy najwyzszej potedze
            try:
              tmp_coeff = modDiv(max_coeff_p1, divisor_coeff, m)
            except ZeroDivisionError as e:
                raise e
            
            tmp_exp = len(p1.poly)-1 - divisor_exp
            tmp = [0] * tmp_exp
            tmp.append(tmp_coeff)
            sub = Polynomial(tmp, m) * p2
            p1 = p1 - sub
        p1.normalize()
        return Polynomial(p1.poly, m)

    def poly_gcd(self, p2):
        p1 = self
        try:
            divisible = p2
        except ZeroDivisionError as e:
            raise e
        if p2.poly == []:
            return p1
        
        return p2.poly_gcd(p1 % p2)

def modDiv(a, b, m): # a*b^-1 (mod m)
    if gcd(b, m) != 1:
        raise ZeroDivisionError
    else:
        return (a * modinv(b, m)) % m
#rozszerzony algorytm euklidesa
def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)

def modinv(a, m):
    g, x, y = egcd(a, m)
    return x % m