sympycore/arithmetic/number_theory.py

   1 """Provides algorithms from number theory.
   2 """
   3
   4 from ..core import init_module
   5
   6 init_module.import_lowlevel_operations()
   7
   8 from .numbers import mpq, div
   9 from . import mpmath
  10
  11 __all__ = ['gcd', 'lcm', 'factorial',
  12            'integer_digits', 'real_digits',
  13            'multinomial_coefficients', 'f2q']
  14
  15 __docformat__ = "restructuredtext en"
  16
  17 def factorial(n, memo=[1, 1]):
  18     """Return n factorial (for integers n >= 0 only)."""
  19     if n < 0:
  20         raise ValueError("expected non-negative integer, got %r" % (n)) #pragma NO COVER
  21     k = len(memo)
  22     if n < k:
  23         return memo[n]
  24     p = memo[-1]
  25     while k <= n:
  26         p *= k
  27         k += 1
  28         if k < 100:
  29             memo.append(p)
  30     return p
  31
  32 def gcd(*args):
  33     """Calculate the greatest common divisor (GCD) of the arguments."""
  34     L = len(args)
  35     if L == 0: return 0
  36     if L == 1: return args[0]
  37     if L == 2:
  38         a, b = args
  39         while b:
  40             a, b = b, a % b
  41         return a
  42     return gcd(gcd(args[0], args[1]), *args[2:])
  43
  44 def lcm(*args):
  45     """Calculate the least common multiple (LCM) of the arguments."""
  46     L = len(args)
  47     if L == 0: return 0
  48     if L == 1: return args[0]
  49     if L == 2: return div(args[0], gcd(*args))*args[1]
  50     return lcm(lcm(args[0], args[1]), *args[2:])
  51
  52 # TODO: this could use the faster implementation in mpmath
  53 def integer_digits(n, base=10):
  54     """Return a list of the digits of abs(n) in the given base."""
  55     assert base > 1
  56     assert isinstance(n, (int, long))
  57     n = abs(n)
  58     if not n:
  59         return [0]
  60     L = []
  61     while n:
  62         n, digit = divmod(n, base)
  63         L.append(int(digit))
  64     return L[::-1]
  65
  66 # TODO: this could (also?) be implemented as an endless generator
  67 def real_digits(x, base=10, truncation=10):
  68     """Return ``(L, d)`` where L is a list of digits of ``abs(x)`` in
  69     the given base and ``d`` is the (signed) distance from the leading
  70     digit to the radix point.
  71
  72     For example, 1234.56 becomes ``([1, 2, 3, 4, 5, 6], 4)`` and 0.001
  73     becomes ``([1], -2)``. If, during the generation of fractional
  74     digits, the length reaches `truncation` digits, the iteration is
  75     stopped."""
  76     assert base > 1
  77     assert isinstance(x, (int, long, mpq))
  78     if x == 0:
  79         return ([0], 1)
  80     x = abs(x)
  81     exponent = 0
  82     while x < 1:
  83         x *= base
  84         exponent -= 1
  85     integer, fraction = divmod(x, 1)
  86     L = integer_digits(integer, base)
  87     exponent += len(L)
  88     if fraction:
  89         p, q = fraction
  90         for i in xrange(truncation - len(L)):
  91             p = (p % q) * base
  92             if not p:
  93                 break
  94             L.append(int(p//q))
  95     return L, exponent
  96
  97 def binomial_coefficients(n):
  98     """Return a dictionary containing pairs {(k1,k2) : C_kn} where
  99     C_kn are binomial coefficients and n=k1+k2.
 100
 101     INPUT:
 102         n -- an integer
 103
 104     OUTPUT:
 105         dict
 106
 107     EXAMPLES:
 108         >>> sorted(binomial_coefficients(3).items())
 109         [((0, 3), 1), ((1, 2), 3), ((2, 1), 3), ((3, 0), 1)]
 110
 111     Notice the coefficients above are the same as below:
 112         x**3 + 3*x**2*y + 3*x*y**2 + y^3
 113
 114     """
 115     d = {(0, n):1, (n, 0):1}
 116     a = 1
 117     for k in xrange(1, n//2+1):
 118         a = (a * (n-k+1))//k
 119         d[k, n-k] = d[n-k, k] = a
 120     return d
 121
 122 def binomial_coefficients_list(n):
 123     """Return a list of binomial coefficients as rows of the Pascal's
 124     triangle.
 125     """
 126     d = [1] * (n+1)
 127     a = 1
 128     for k in xrange(1, n//2+1):
 129         a = (a * (n-k+1))//k
 130         d[k] = d[n-k] = a
 131     return d
 132
 133 def multinomial_coefficients(m, n, _tuple=tuple, _zip=zip):
 134     """Return a dictionary containing pairs ``{(k1,k2,..,km) : C_kn}``
 135     where ``C_kn`` are multinomial coefficients such that
 136     ``n=k1+k2+..+km``.
 137
 138     INPUT:
 139         m -- integer
 140         n -- integer
 141         _tuple, _zip -- hacks for speed; don't set these as a user.
 142
 143     OUTPUT:
 144         dict
 145
 146     EXAMPLES:
 147         >>> sorted(multinomial_coefficients(2,5).items())
 148         [((0, 5), 1), ((1, 4), 5), ((2, 3), 10), ((3, 2), 10), ((4, 1), 5), ((5, 0), 1)]
 149
 150     Notice that these are the coefficients of ``(x+y)**5``:
 151          x**5 + 5*x**4*y + 10*x**3*y**2 + 10*x**2*y**3 + 5*x*y**4 + y**5
 152
 153         >>> sorted(multinomial_coefficients(3,2).items())
 154         [((0, 0, 2), 1), ((0, 1, 1), 2), ((0, 2, 0), 1), ((1, 0, 1), 2), ((1, 1, 0), 2), ((2, 0, 0), 1)]
 155
 156
 157     ALGORITHM:
 158     The algorithm we implement for computing the multinomial coefficients
 159     is based on the following result:
 160
 161        Consider a polynomial and its ``n``-th exponent::
 162
 163          P(x) = sum_{i=0}^m p_i x^k
 164          P(x)^n = sum_{k=0}^{m n} a(n,k) x^k
 165
 166        We compute the coefficients ``a(n,k)`` using the
 167        J.C.P. Miller Pure Recurrence [see D.E.Knuth, Seminumerical
 168        Algorithms, The art of Computer Programming v.2, Addison
 169        Wesley, Reading, 1981;]::
 170
 171          a(n,k) = 1/(k p_0) sum_{i=1}^m p_i ((n+1)i-k) a(n,k-i),
 172
 173        where ``a(n,0) = p_0^n``.
 174     """
 175     if m==0:
 176         return {}
 177     if m==2:
 178         return binomial_coefficients(n)
 179     symbols = [(0,)*i + (1,) + (0,)*(m-i-1) for i in range(m)]
 180     s0 = symbols[0]
 181     p0 = [_tuple([aa-bb for aa,bb in _zip(s,s0)]) for s in symbols]
 182     r = {_tuple([aa*n for aa in s0]):1}
 183     r_get = r.get
 184     r_update = r.update
 185     l = [0] * (n*(m-1)+1)
 186     l[0] = r.items()
 187     for k in xrange(1, n*(m-1)+1):
 188         d = {}
 189         d_get = d.get
 190         for i in xrange(1, min(m,k+1)):
 191             nn = (n+1)*i-k
 192             if not nn:
 193                 continue
 194             t = p0[i]
 195             for t2, c2 in l[k-i]:
 196                 tt = _tuple([aa+bb for aa,bb in _zip(t2,t)])
 197                 cc = nn * c2
 198                 dict_add_item(None, d, tt, cc)
 199         r1 = [(t, c//k) for (t, c) in d.iteritems()]
 200         l[k] = r1
 201         r_update(r1)
 202     return r
 203
 204 def reldiff(x,y):
 205     """abs(x-y)/((abs(x)+abs(y))/2)"""
 206     return abs(x-y)/((abs(x)+abs(y))/2)
 207
 208 def f2q(f, minerr=None):
 209     '''
 210     Find (almost) the best rational representation of a float that has
 211     finite precision.
 212
 213     The algorithm is based on Stern-Brocot representation of
 214     irrational/rational numbers and its relation to continuants [1].
 215
 216     References:
 217     [1] "Concrete Mathematics" by Graham, Knuth, and Patashnik. Addison-Wesley. 1989
 218     '''
 219     if mpmath.libmp.BACKEND=='gmpy':
 220         tmpz = type(mpmath.libmp.MPZ_ZERO)
 221         if minerr is None:
 222             r = mpmath.libmp.gmpy.f2q(f)
 223         else:
 224             r = mpmath.libmp.gmpy.f2q(f, minerr)
 225         if isinstance (r, tmpz):
 226             return int(r)
 227         n,d = int(r.numer ()), int(r.denom ())
 228         if d==1:
 229             return n
 230         return mpq((n,d))
 231     mpf = mpmath.mpf
 232     f = mpf(f)
 233     if f<0: return -f2q(-f,minerr)
 234     al = f                          # Initialize Eq.(6.142) [1]
 235     if not minerr:
 236         minerr = 1/mpf(2**al.context.prec) # min number for the given precision
 237     elif minerr<0:
 238         minerr = 1/mpf(2**-minerr)
 239     a = mpmath.floor(al)   # a = floor(al) # Initialize Eq.(6.142) [1]
 240     r1=[0,0,1]                           # r1[1:] is 1st row of Eq.(6.138) [1]
 241     r2=[0,1,a]                           # r2[1:] is 2nd row of Eq.(6.138) [1]
 242     err = reldiff(f,a)
 243     while err>minerr:
 244         al = 1/(al-a)                    # Eq.(6.142) [1]
 245         a = mpmath.floor(al) # a = floor(al) # Eq.(6.142) [1]
 246         r1[0] = r1[1]                    # roll
 247         r1[1] = r1[2]                    # roll
 248         r1[2] = r1[1]*a+r1[0]            # use Eq.(6.127) [1]
 249         r2[0] = r2[1]                    # roll
 250         r2[1] = r2[2]                    # roll
 251         r2[2] = r2[1]*a+r2[0]            # use Eq.(6.127) [1]
 252         newerr = reldiff(f,r2[2]/r1[2])
 253         if err<=newerr:
 254             n, d = (int(r2[1]),int(r1[1]))
 255             if d==1: return n
 256             return mpq((n, d))
 257         err = newerr
 258     n, d = (int(r2[2]),int(r1[2]))
 259     if d==1: return n
 260     return mpq((n, d))
 261
 262
 263 #from .combinatorics import multinomial_coefficients