hensel.py

@@ -0,0 +1,29 @@
+# Finding roots of polynomials in p-adic integers using Hensel's lemma
+
+from padic import *
+from poly import *
+
+class PAdicPoly(PAdic):
+    'Result of lifting a root of a polynomial in the integers mod p to the p-adic integers.'
+    def __init__(self, p, poly, root):
+        PAdic.__init__(self, p)
+        self.root = root
+        self.poly = poly
+        self.deriv = derivative(poly)
+        
+        # argument checks for the algorithm to work
+        if poly(root) % p:
+            raise ValueError("%d is not a root of %s modulo %d" % (root, poly, p))
+        if self.deriv(root) % p == 0:
+            raise ValueError("Polynomial %s is not separable modulo %d" % (poly, p))
+        
+        # take care of trailing zeros
+        digit = self.root
+        while digit == 0:
+            digit = self._nextdigit()
+            self.order += 1
+            self.prec += 1
+
+    def _nextdigit(self):
+        self.root = ModP(self.p ** (self.order + 2), self.root)
+        self.root = self.root - self.poly(self.root) / self.deriv(self.root) # coercions automatically taken care of

modp.py

@@ -2,7 +2,7 @@ class ModP(int):
     'Integers mod p, p a prime power.'
     def __new__(cls, p, num):
         self.p = p
-        return int.__new__(cls, num % p)
+        return int.__new__(cls, int(num) % p)
     
     # arithmetic
     def __add__(self, other):

poly.py

@@ -8,6 +8,19 @@ class Poly:
         self.coeffs = defaultdict(int, not coeffs and {} or isinstance(coeffs, int) and {0:coeffs} or coeffs)
         self.deg = int(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):
+            return ' * '.join(([] if coeff == 1 else [str(coeff)]) + ([] if expt == 0 else ['X'] if expt == 1 else ['X^' + expt]))
+            
+        return ' + '.join(term(self.coeffs[i], i) for i in self.coeffs if self.coeffs[i] != 0)
+    
     # arithmetic
     def __add__(self, other):
         if not isinstance(other, Poly):
@@ -34,7 +47,7 @@ class Poly:
         if not isinstance(other, Poly):
             other = Poly(other)
         res = Poly()
-        res.deg = self.deg + other.deg
+        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]
@@ -44,4 +57,8 @@ class Poly:
             other = Poly(other)
         return other.__mul__(self)
 
-X = Poly({1:1})
+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})