padic.py

@@ -6,15 +6,12 @@ class PAdic:
     def __init__(self, p):
         self.p = p
         self.val  = '' # current known value
-        self.prec = 0 # current known precision
+        self.prec = 0 # current known precision, also containing leading zeros
         self.order = 0 # order/valuation of number
         pass # initialize object as subclass perhaps
     
     def get(self, prec):
         'Return value of number with given precision.'
-        if self.value == 0:
-            return '0' # * prec
-        
         while self.prec < prec:
             # update val based on value
             self.val = self._nextdigit() + self.val
@@ -71,20 +68,24 @@ class PAdicConst(PAdic):
         if value == 0:
             self.value = value
             self.val = '0'
+            self.order = maxint
             return
         
-        self.norm = Fraction(1)
         self.order = 0
         while not value.numerator % self.p:
-            self.norm /= self.p
             self.order += 1
+            self.prec += 1
             value.numerator /= self.p
         while not value.denominator % self.p:
-            self.norm *= self.p
             valuation -= 1
             value.denominator /= self.p
         self.value = value
     
+    def get(self, prec):
+        if self.value == 0:
+            return '0' # * prec
+        return PAdic.get(self, prec)
+    
     def _nextdigit(self):
         'Calculate next digit of p-adic number.'
         rem = ModP(self.p, self.value.numerator) / ModP(self.p, self.value.denominator)
@@ -99,6 +100,13 @@ class PAdicAdd(PAdic):
         self.carry = 0
         self.arg1 = arg1
         self.arg2 = arg2
+        self.order = min(arg1.order, arg2.order) # might be larger than this
+        # loop until first nonzero digit is found
+        digit = self._nextdigit()
+        while digit == 0:
+            self.order += 1
+            self.prec += 1
+            digit = self._nextdigit()
     
     def _nextdigit(self):
         s = self.arg1.getdigit(self.prec) + self.arg2.getdigit(self.prec) + self.carry
@@ -110,10 +118,11 @@ class PAdicNeg(PAdic):
     def __init__(self, p, arg):
         PAdic.__init__(self, p)
         self.arg = arg
+        self.order = arg.order
     
     def _nextdigit(self):
         if self.prec == 0:
-            return self.p - self.arg.getdigit(0)
+            return self.p - self.arg.getdigit(0) # cannot be p, 0th digit of arg must be nonzero
         return self.p - 1 - self.arg.getdigit(0)
 
 class PAdicMul(PAdic):

poly.py

@@ -1,18 +1,47 @@
-# Finding roots of polynomials in p-adic integers using Hensel's lemma
+# Polynomial class on Z or Z_p
 
-from padic import *
 from collections import defaultdict
 
 class Poly:
     'Polynomial class.'
     def __init__(self, coeffs = None):
-        self.coeffs = defaultdict(int, coeffs or {})
-        self.deg = 0
-    # TODO arithmetic
+        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()))
+    
+    # arithmetic
+    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
+        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)
 
-X = Poly({1:1})
-
-class ConstPoly(Poly):
-    'Constant polynomial.'
-    def __init__(self, coeff):
-        Poly.__init__(self, {0:coeff})
+X = Poly({1:1})