Coverage for /usr/lib/python3/dist-packages/mpmath/libmp/libintmath.py: 31%
345 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2Utility functions for integer math.
4TODO: rename, cleanup, perhaps move the gmpy wrapper code
5here from settings.py
7"""
9import math
10from bisect import bisect
12from .backend import xrange
13from .backend import BACKEND, gmpy, sage, sage_utils, MPZ, MPZ_ONE, MPZ_ZERO
15small_trailing = [0] * 256
16for j in range(1,8):
17 small_trailing[1<<j::1<<(j+1)] = [j] * (1<<(7-j))
19def giant_steps(start, target, n=2):
20 """
21 Return a list of integers ~=
23 [start, n*start, ..., target/n^2, target/n, target]
25 but conservatively rounded so that the quotient between two
26 successive elements is actually slightly less than n.
28 With n = 2, this describes suitable precision steps for a
29 quadratically convergent algorithm such as Newton's method;
30 with n = 3 steps for cubic convergence (Halley's method), etc.
32 >>> giant_steps(50,1000)
33 [66, 128, 253, 502, 1000]
34 >>> giant_steps(50,1000,4)
35 [65, 252, 1000]
37 """
38 L = [target]
39 while L[-1] > start*n:
40 L = L + [L[-1]//n + 2]
41 return L[::-1]
43def rshift(x, n):
44 """For an integer x, calculate x >> n with the fastest (floor)
45 rounding. Unlike the plain Python expression (x >> n), n is
46 allowed to be negative, in which case a left shift is performed."""
47 if n >= 0: return x >> n
48 else: return x << (-n)
50def lshift(x, n):
51 """For an integer x, calculate x << n. Unlike the plain Python
52 expression (x << n), n is allowed to be negative, in which case a
53 right shift with default (floor) rounding is performed."""
54 if n >= 0: return x << n
55 else: return x >> (-n)
57if BACKEND == 'sage':
58 import operator
59 rshift = operator.rshift
60 lshift = operator.lshift
62def python_trailing(n):
63 """Count the number of trailing zero bits in abs(n)."""
64 if not n:
65 return 0
66 low_byte = n & 0xff
67 if low_byte:
68 return small_trailing[low_byte]
69 t = 8
70 n >>= 8
71 while not n & 0xff:
72 n >>= 8
73 t += 8
74 return t + small_trailing[n & 0xff]
76if BACKEND == 'gmpy':
77 if gmpy.version() >= '2':
78 def gmpy_trailing(n):
79 """Count the number of trailing zero bits in abs(n) using gmpy."""
80 if n: return MPZ(n).bit_scan1()
81 else: return 0
82 else:
83 def gmpy_trailing(n):
84 """Count the number of trailing zero bits in abs(n) using gmpy."""
85 if n: return MPZ(n).scan1()
86 else: return 0
88# Small powers of 2
89powers = [1<<_ for _ in range(300)]
91def python_bitcount(n):
92 """Calculate bit size of the nonnegative integer n."""
93 bc = bisect(powers, n)
94 if bc != 300:
95 return bc
96 bc = int(math.log(n, 2)) - 4
97 return bc + bctable[n>>bc]
99def gmpy_bitcount(n):
100 """Calculate bit size of the nonnegative integer n."""
101 if n: return MPZ(n).numdigits(2)
102 else: return 0
104#def sage_bitcount(n):
105# if n: return MPZ(n).nbits()
106# else: return 0
108def sage_trailing(n):
109 return MPZ(n).trailing_zero_bits()
111if BACKEND == 'gmpy':
112 bitcount = gmpy_bitcount
113 trailing = gmpy_trailing
114elif BACKEND == 'sage':
115 sage_bitcount = sage_utils.bitcount
116 bitcount = sage_bitcount
117 trailing = sage_trailing
118else:
119 bitcount = python_bitcount
120 trailing = python_trailing
122if BACKEND == 'gmpy' and 'bit_length' in dir(gmpy):
123 bitcount = gmpy.bit_length
125# Used to avoid slow function calls as far as possible
126trailtable = [trailing(n) for n in range(256)]
127bctable = [bitcount(n) for n in range(1024)]
129# TODO: speed up for bases 2, 4, 8, 16, ...
131def bin_to_radix(x, xbits, base, bdigits):
132 """Changes radix of a fixed-point number; i.e., converts
133 x * 2**xbits to floor(x * 10**bdigits)."""
134 return x * (MPZ(base)**bdigits) >> xbits
136stddigits = '0123456789abcdefghijklmnopqrstuvwxyz'
138def small_numeral(n, base=10, digits=stddigits):
139 """Return the string numeral of a positive integer in an arbitrary
140 base. Most efficient for small input."""
141 if base == 10:
142 return str(n)
143 digs = []
144 while n:
145 n, digit = divmod(n, base)
146 digs.append(digits[digit])
147 return "".join(digs[::-1])
149def numeral_python(n, base=10, size=0, digits=stddigits):
150 """Represent the integer n as a string of digits in the given base.
151 Recursive division is used to make this function about 3x faster
152 than Python's str() for converting integers to decimal strings.
154 The 'size' parameters specifies the number of digits in n; this
155 number is only used to determine splitting points and need not be
156 exact."""
157 if n <= 0:
158 if not n:
159 return "0"
160 return "-" + numeral(-n, base, size, digits)
161 # Fast enough to do directly
162 if size < 250:
163 return small_numeral(n, base, digits)
164 # Divide in half
165 half = (size // 2) + (size & 1)
166 A, B = divmod(n, base**half)
167 ad = numeral(A, base, half, digits)
168 bd = numeral(B, base, half, digits).rjust(half, "0")
169 return ad + bd
171def numeral_gmpy(n, base=10, size=0, digits=stddigits):
172 """Represent the integer n as a string of digits in the given base.
173 Recursive division is used to make this function about 3x faster
174 than Python's str() for converting integers to decimal strings.
176 The 'size' parameters specifies the number of digits in n; this
177 number is only used to determine splitting points and need not be
178 exact."""
179 if n < 0:
180 return "-" + numeral(-n, base, size, digits)
181 # gmpy.digits() may cause a segmentation fault when trying to convert
182 # extremely large values to a string. The size limit may need to be
183 # adjusted on some platforms, but 1500000 works on Windows and Linux.
184 if size < 1500000:
185 return gmpy.digits(n, base)
186 # Divide in half
187 half = (size // 2) + (size & 1)
188 A, B = divmod(n, MPZ(base)**half)
189 ad = numeral(A, base, half, digits)
190 bd = numeral(B, base, half, digits).rjust(half, "0")
191 return ad + bd
193if BACKEND == "gmpy":
194 numeral = numeral_gmpy
195else:
196 numeral = numeral_python
198_1_800 = 1<<800
199_1_600 = 1<<600
200_1_400 = 1<<400
201_1_200 = 1<<200
202_1_100 = 1<<100
203_1_50 = 1<<50
205def isqrt_small_python(x):
206 """
207 Correctly (floor) rounded integer square root, using
208 division. Fast up to ~200 digits.
209 """
210 if not x:
211 return x
212 if x < _1_800:
213 # Exact with IEEE double precision arithmetic
214 if x < _1_50:
215 return int(x**0.5)
216 # Initial estimate can be any integer >= the true root; round up
217 r = int(x**0.5 * 1.00000000000001) + 1
218 else:
219 bc = bitcount(x)
220 n = bc//2
221 r = int((x>>(2*n-100))**0.5+2)<<(n-50) # +2 is to round up
222 # The following iteration now precisely computes floor(sqrt(x))
223 # See e.g. Crandall & Pomerance, "Prime Numbers: A Computational
224 # Perspective"
225 while 1:
226 y = (r+x//r)>>1
227 if y >= r:
228 return r
229 r = y
231def isqrt_fast_python(x):
232 """
233 Fast approximate integer square root, computed using division-free
234 Newton iteration for large x. For random integers the result is almost
235 always correct (floor(sqrt(x))), but is 1 ulp too small with a roughly
236 0.1% probability. If x is very close to an exact square, the answer is
237 1 ulp wrong with high probability.
239 With 0 guard bits, the largest error over a set of 10^5 random
240 inputs of size 1-10^5 bits was 3 ulp. The use of 10 guard bits
241 almost certainly guarantees a max 1 ulp error.
242 """
243 # Use direct division-based iteration if sqrt(x) < 2^400
244 # Assume floating-point square root accurate to within 1 ulp, then:
245 # 0 Newton iterations good to 52 bits
246 # 1 Newton iterations good to 104 bits
247 # 2 Newton iterations good to 208 bits
248 # 3 Newton iterations good to 416 bits
249 if x < _1_800:
250 y = int(x**0.5)
251 if x >= _1_100:
252 y = (y + x//y) >> 1
253 if x >= _1_200:
254 y = (y + x//y) >> 1
255 if x >= _1_400:
256 y = (y + x//y) >> 1
257 return y
258 bc = bitcount(x)
259 guard_bits = 10
260 x <<= 2*guard_bits
261 bc += 2*guard_bits
262 bc += (bc&1)
263 hbc = bc//2
264 startprec = min(50, hbc)
265 # Newton iteration for 1/sqrt(x), with floating-point starting value
266 r = int(2.0**(2*startprec) * (x >> (bc-2*startprec)) ** -0.5)
267 pp = startprec
268 for p in giant_steps(startprec, hbc):
269 # r**2, scaled from real size 2**(-bc) to 2**p
270 r2 = (r*r) >> (2*pp - p)
271 # x*r**2, scaled from real size ~1.0 to 2**p
272 xr2 = ((x >> (bc-p)) * r2) >> p
273 # New value of r, scaled from real size 2**(-bc/2) to 2**p
274 r = (r * ((3<<p) - xr2)) >> (pp+1)
275 pp = p
276 # (1/sqrt(x))*x = sqrt(x)
277 return (r*(x>>hbc)) >> (p+guard_bits)
279def sqrtrem_python(x):
280 """Correctly rounded integer (floor) square root with remainder."""
281 # to check cutoff:
282 # plot(lambda x: timing(isqrt, 2**int(x)), [0,2000])
283 if x < _1_600:
284 y = isqrt_small_python(x)
285 return y, x - y*y
286 y = isqrt_fast_python(x) + 1
287 rem = x - y*y
288 # Correct remainder
289 while rem < 0:
290 y -= 1
291 rem += (1+2*y)
292 else:
293 if rem:
294 while rem > 2*(1+y):
295 y += 1
296 rem -= (1+2*y)
297 return y, rem
299def isqrt_python(x):
300 """Integer square root with correct (floor) rounding."""
301 return sqrtrem_python(x)[0]
303def sqrt_fixed(x, prec):
304 return isqrt_fast(x<<prec)
306sqrt_fixed2 = sqrt_fixed
308if BACKEND == 'gmpy':
309 if gmpy.version() >= '2':
310 isqrt_small = isqrt_fast = isqrt = gmpy.isqrt
311 sqrtrem = gmpy.isqrt_rem
312 else:
313 isqrt_small = isqrt_fast = isqrt = gmpy.sqrt
314 sqrtrem = gmpy.sqrtrem
315elif BACKEND == 'sage':
316 isqrt_small = isqrt_fast = isqrt = \
317 getattr(sage_utils, "isqrt", lambda n: MPZ(n).isqrt())
318 sqrtrem = lambda n: MPZ(n).sqrtrem()
319else:
320 isqrt_small = isqrt_small_python
321 isqrt_fast = isqrt_fast_python
322 isqrt = isqrt_python
323 sqrtrem = sqrtrem_python
326def ifib(n, _cache={}):
327 """Computes the nth Fibonacci number as an integer, for
328 integer n."""
329 if n < 0:
330 return (-1)**(-n+1) * ifib(-n)
331 if n in _cache:
332 return _cache[n]
333 m = n
334 # Use Dijkstra's logarithmic algorithm
335 # The following implementation is basically equivalent to
336 # http://en.literateprograms.org/Fibonacci_numbers_(Scheme)
337 a, b, p, q = MPZ_ONE, MPZ_ZERO, MPZ_ZERO, MPZ_ONE
338 while n:
339 if n & 1:
340 aq = a*q
341 a, b = b*q+aq+a*p, b*p+aq
342 n -= 1
343 else:
344 qq = q*q
345 p, q = p*p+qq, qq+2*p*q
346 n >>= 1
347 if m < 250:
348 _cache[m] = b
349 return b
351MAX_FACTORIAL_CACHE = 1000
353def ifac(n, memo={0:1, 1:1}):
354 """Return n factorial (for integers n >= 0 only)."""
355 f = memo.get(n)
356 if f:
357 return f
358 k = len(memo)
359 p = memo[k-1]
360 MAX = MAX_FACTORIAL_CACHE
361 while k <= n:
362 p *= k
363 if k <= MAX:
364 memo[k] = p
365 k += 1
366 return p
368def ifac2(n, memo_pair=[{0:1}, {1:1}]):
369 """Return n!! (double factorial), integers n >= 0 only."""
370 memo = memo_pair[n&1]
371 f = memo.get(n)
372 if f:
373 return f
374 k = max(memo)
375 p = memo[k]
376 MAX = MAX_FACTORIAL_CACHE
377 while k < n:
378 k += 2
379 p *= k
380 if k <= MAX:
381 memo[k] = p
382 return p
384if BACKEND == 'gmpy':
385 ifac = gmpy.fac
386elif BACKEND == 'sage':
387 ifac = lambda n: int(sage.factorial(n))
388 ifib = sage.fibonacci
390def list_primes(n):
391 n = n + 1
392 sieve = list(xrange(n))
393 sieve[:2] = [0, 0]
394 for i in xrange(2, int(n**0.5)+1):
395 if sieve[i]:
396 for j in xrange(i**2, n, i):
397 sieve[j] = 0
398 return [p for p in sieve if p]
400if BACKEND == 'sage':
401 # Note: it is *VERY* important for performance that we convert
402 # the list to Python ints.
403 def list_primes(n):
404 return [int(_) for _ in sage.primes(n+1)]
406small_odd_primes = (3,5,7,11,13,17,19,23,29,31,37,41,43,47)
407small_odd_primes_set = set(small_odd_primes)
409def isprime(n):
410 """
411 Determines whether n is a prime number. A probabilistic test is
412 performed if n is very large. No special trick is used for detecting
413 perfect powers.
415 >>> sum(list_primes(100000))
416 454396537
417 >>> sum(n*isprime(n) for n in range(100000))
418 454396537
420 """
421 n = int(n)
422 if not n & 1:
423 return n == 2
424 if n < 50:
425 return n in small_odd_primes_set
426 for p in small_odd_primes:
427 if not n % p:
428 return False
429 m = n-1
430 s = trailing(m)
431 d = m >> s
432 def test(a):
433 x = pow(a,d,n)
434 if x == 1 or x == m:
435 return True
436 for r in xrange(1,s):
437 x = x**2 % n
438 if x == m:
439 return True
440 return False
441 # See http://primes.utm.edu/prove/prove2_3.html
442 if n < 1373653:
443 witnesses = [2,3]
444 elif n < 341550071728321:
445 witnesses = [2,3,5,7,11,13,17]
446 else:
447 witnesses = small_odd_primes
448 for a in witnesses:
449 if not test(a):
450 return False
451 return True
453def moebius(n):
454 """
455 Evaluates the Moebius function which is `mu(n) = (-1)^k` if `n`
456 is a product of `k` distinct primes and `mu(n) = 0` otherwise.
458 TODO: speed up using factorization
459 """
460 n = abs(int(n))
461 if n < 2:
462 return n
463 factors = []
464 for p in xrange(2, n+1):
465 if not (n % p):
466 if not (n % p**2):
467 return 0
468 if not sum(p % f for f in factors):
469 factors.append(p)
470 return (-1)**len(factors)
472def gcd(*args):
473 a = 0
474 for b in args:
475 if a:
476 while b:
477 a, b = b, a % b
478 else:
479 a = b
480 return a
483# Comment by Juan Arias de Reyna:
484#
485# I learn this method to compute EulerE[2n] from van de Lune.
486#
487# We apply the formula EulerE[2n] = (-1)^n 2**(-2n) sum_{j=0}^n a(2n,2j+1)
488#
489# where the numbers a(n,j) vanish for j > n+1 or j <= -1 and satisfies
490#
491# a(0,-1) = a(0,0) = 0; a(0,1)= 1; a(0,2) = a(0,3) = 0
492#
493# a(n,j) = a(n-1,j) when n+j is even
494# a(n,j) = (j-1) a(n-1,j-1) + (j+1) a(n-1,j+1) when n+j is odd
495#
496#
497# But we can use only one array unidimensional a(j) since to compute
498# a(n,j) we only need to know a(n-1,k) where k and j are of different parity
499# and we have not to conserve the used values.
500#
501# We cached up the values of Euler numbers to sufficiently high order.
502#
503# Important Observation: If we pretend to use the numbers
504# EulerE[1], EulerE[2], ... , EulerE[n]
505# it is convenient to compute first EulerE[n], since the algorithm
506# computes first all
507# the previous ones, and keeps them in the CACHE
509MAX_EULER_CACHE = 500
511def eulernum(m, _cache={0:MPZ_ONE}):
512 r"""
513 Computes the Euler numbers `E(n)`, which can be defined as
514 coefficients of the Taylor expansion of `1/cosh x`:
516 .. math ::
518 \frac{1}{\cosh x} = \sum_{n=0}^\infty \frac{E_n}{n!} x^n
520 Example::
522 >>> [int(eulernum(n)) for n in range(11)]
523 [1, 0, -1, 0, 5, 0, -61, 0, 1385, 0, -50521]
524 >>> [int(eulernum(n)) for n in range(11)] # test cache
525 [1, 0, -1, 0, 5, 0, -61, 0, 1385, 0, -50521]
527 """
528 # for odd m > 1, the Euler numbers are zero
529 if m & 1:
530 return MPZ_ZERO
531 f = _cache.get(m)
532 if f:
533 return f
534 MAX = MAX_EULER_CACHE
535 n = m
536 a = [MPZ(_) for _ in [0,0,1,0,0,0]]
537 for n in range(1, m+1):
538 for j in range(n+1, -1, -2):
539 a[j+1] = (j-1)*a[j] + (j+1)*a[j+2]
540 a.append(0)
541 suma = 0
542 for k in range(n+1, -1, -2):
543 suma += a[k+1]
544 if n <= MAX:
545 _cache[n] = ((-1)**(n//2))*(suma // 2**n)
546 if n == m:
547 return ((-1)**(n//2))*suma // 2**n
549def stirling1(n, k):
550 """
551 Stirling number of the first kind.
552 """
553 if n < 0 or k < 0:
554 raise ValueError
555 if k >= n:
556 return MPZ(n == k)
557 if k < 1:
558 return MPZ_ZERO
559 L = [MPZ_ZERO] * (k+1)
560 L[1] = MPZ_ONE
561 for m in xrange(2, n+1):
562 for j in xrange(min(k, m), 0, -1):
563 L[j] = (m-1) * L[j] + L[j-1]
564 return (-1)**(n+k) * L[k]
566def stirling2(n, k):
567 """
568 Stirling number of the second kind.
569 """
570 if n < 0 or k < 0:
571 raise ValueError
572 if k >= n:
573 return MPZ(n == k)
574 if k <= 1:
575 return MPZ(k == 1)
576 s = MPZ_ZERO
577 t = MPZ_ONE
578 for j in xrange(k+1):
579 if (k + j) & 1:
580 s -= t * MPZ(j)**n
581 else:
582 s += t * MPZ(j)**n
583 t = t * (k - j) // (j + 1)
584 return s // ifac(k)