ntheory.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618
  1. # sympy.external.ntheory
  2. #
  3. # This module provides pure Python implementations of some number theory
  4. # functions that are alternately used from gmpy2 if it is installed.
  5. import math
  6. import mpmath.libmp as mlib
  7. _small_trailing = [0] * 256
  8. for j in range(1, 8):
  9. _small_trailing[1 << j :: 1 << (j + 1)] = [j] * (1 << (7 - j))
  10. def bit_scan1(x, n=0):
  11. if not x:
  12. return
  13. x = abs(x >> n)
  14. low_byte = x & 0xFF
  15. if low_byte:
  16. return _small_trailing[low_byte] + n
  17. t = 8 + n
  18. x >>= 8
  19. # 2**m is quick for z up through 2**30
  20. z = x.bit_length() - 1
  21. if x == 1 << z:
  22. return z + t
  23. if z < 300:
  24. # fixed 8-byte reduction
  25. while not x & 0xFF:
  26. x >>= 8
  27. t += 8
  28. else:
  29. # binary reduction important when there might be a large
  30. # number of trailing 0s
  31. p = z >> 1
  32. while not x & 0xFF:
  33. while x & ((1 << p) - 1):
  34. p >>= 1
  35. x >>= p
  36. t += p
  37. return t + _small_trailing[x & 0xFF]
  38. def bit_scan0(x, n=0):
  39. return bit_scan1(x + (1 << n), n)
  40. def remove(x, f):
  41. if f < 2:
  42. raise ValueError("factor must be > 1")
  43. if x == 0:
  44. return 0, 0
  45. if f == 2:
  46. b = bit_scan1(x)
  47. return x >> b, b
  48. m = 0
  49. y, rem = divmod(x, f)
  50. while not rem:
  51. x = y
  52. m += 1
  53. if m > 5:
  54. pow_list = [f**2]
  55. while pow_list:
  56. _f = pow_list[-1]
  57. y, rem = divmod(x, _f)
  58. if not rem:
  59. m += 1 << len(pow_list)
  60. x = y
  61. pow_list.append(_f**2)
  62. else:
  63. pow_list.pop()
  64. y, rem = divmod(x, f)
  65. return x, m
  66. def factorial(x):
  67. """Return x!."""
  68. return int(mlib.ifac(int(x)))
  69. def sqrt(x):
  70. """Integer square root of x."""
  71. return int(mlib.isqrt(int(x)))
  72. def sqrtrem(x):
  73. """Integer square root of x and remainder."""
  74. s, r = mlib.sqrtrem(int(x))
  75. return (int(s), int(r))
  76. gcd = math.gcd
  77. lcm = math.lcm
  78. def _sign(n):
  79. if n < 0:
  80. return -1, -n
  81. return 1, n
  82. def gcdext(a, b):
  83. if not a or not b:
  84. g = abs(a) or abs(b)
  85. if not g:
  86. return (0, 0, 0)
  87. return (g, a // g, b // g)
  88. x_sign, a = _sign(a)
  89. y_sign, b = _sign(b)
  90. x, r = 1, 0
  91. y, s = 0, 1
  92. while b:
  93. q, c = divmod(a, b)
  94. a, b = b, c
  95. x, r = r, x - q*r
  96. y, s = s, y - q*s
  97. return (a, x * x_sign, y * y_sign)
  98. def is_square(x):
  99. """Return True if x is a square number."""
  100. if x < 0:
  101. return False
  102. # Note that the possible values of y**2 % n for a given n are limited.
  103. # For example, when n=4, y**2 % n can only take 0 or 1.
  104. # In other words, if x % 4 is 2 or 3, then x is not a square number.
  105. # Mathematically, it determines if it belongs to the set {y**2 % n},
  106. # but implementationally, it can be realized as a logical conjunction
  107. # with an n-bit integer.
  108. # see https://mersenneforum.org/showpost.php?p=110896
  109. # def magic(n):
  110. # s = {y**2 % n for y in range(n)}
  111. # s = set(range(n)) - s
  112. # return sum(1 << bit for bit in s)
  113. # >>> print(hex(magic(128)))
  114. # 0xfdfdfdedfdfdfdecfdfdfdedfdfcfdec
  115. # >>> print(hex(magic(99)))
  116. # 0x5f6f9ffb6fb7ddfcb75befdec
  117. # >>> print(hex(magic(91)))
  118. # 0x6fd1bfcfed5f3679d3ebdec
  119. # >>> print(hex(magic(85)))
  120. # 0xdef9ae771ffe3b9d67dec
  121. if 0xfdfdfdedfdfdfdecfdfdfdedfdfcfdec & (1 << (x & 127)):
  122. return False # e.g. 2, 3
  123. m = x % 765765 # 765765 = 99 * 91 * 85
  124. if 0x5f6f9ffb6fb7ddfcb75befdec & (1 << (m % 99)):
  125. return False # e.g. 17, 68
  126. if 0x6fd1bfcfed5f3679d3ebdec & (1 << (m % 91)):
  127. return False # e.g. 97, 388
  128. if 0xdef9ae771ffe3b9d67dec & (1 << (m % 85)):
  129. return False # e.g. 793, 1408
  130. return mlib.sqrtrem(int(x))[1] == 0
  131. def invert(x, m):
  132. """Modular inverse of x modulo m.
  133. Returns y such that x*y == 1 mod m.
  134. Uses ``math.pow`` but reproduces the behaviour of ``gmpy2.invert``
  135. which raises ZeroDivisionError if no inverse exists.
  136. """
  137. try:
  138. return pow(x, -1, m)
  139. except ValueError:
  140. raise ZeroDivisionError("invert() no inverse exists")
  141. def legendre(x, y):
  142. """Legendre symbol (x / y).
  143. Following the implementation of gmpy2,
  144. the error is raised only when y is an even number.
  145. """
  146. if y <= 0 or not y % 2:
  147. raise ValueError("y should be an odd prime")
  148. x %= y
  149. if not x:
  150. return 0
  151. if pow(x, (y - 1) // 2, y) == 1:
  152. return 1
  153. return -1
  154. def jacobi(x, y):
  155. """Jacobi symbol (x / y)."""
  156. if y <= 0 or not y % 2:
  157. raise ValueError("y should be an odd positive integer")
  158. x %= y
  159. if not x:
  160. return int(y == 1)
  161. if y == 1 or x == 1:
  162. return 1
  163. if gcd(x, y) != 1:
  164. return 0
  165. j = 1
  166. while x != 0:
  167. while x % 2 == 0 and x > 0:
  168. x >>= 1
  169. if y % 8 in [3, 5]:
  170. j = -j
  171. x, y = y, x
  172. if x % 4 == y % 4 == 3:
  173. j = -j
  174. x %= y
  175. return j
  176. def kronecker(x, y):
  177. """Kronecker symbol (x / y)."""
  178. if gcd(x, y) != 1:
  179. return 0
  180. if y == 0:
  181. return 1
  182. sign = -1 if y < 0 and x < 0 else 1
  183. y = abs(y)
  184. s = bit_scan1(y)
  185. y >>= s
  186. if s % 2 and x % 8 in [3, 5]:
  187. sign = -sign
  188. return sign * jacobi(x, y)
  189. def iroot(y, n):
  190. if y < 0:
  191. raise ValueError("y must be nonnegative")
  192. if n < 1:
  193. raise ValueError("n must be positive")
  194. if y in (0, 1):
  195. return y, True
  196. if n == 1:
  197. return y, True
  198. if n == 2:
  199. x, rem = mlib.sqrtrem(y)
  200. return int(x), not rem
  201. if n >= y.bit_length():
  202. return 1, False
  203. # Get initial estimate for Newton's method. Care must be taken to
  204. # avoid overflow
  205. try:
  206. guess = int(y**(1./n) + 0.5)
  207. except OverflowError:
  208. exp = math.log2(y)/n
  209. if exp > 53:
  210. shift = int(exp - 53)
  211. guess = int(2.0**(exp - shift) + 1) << shift
  212. else:
  213. guess = int(2.0**exp)
  214. if guess > 2**50:
  215. # Newton iteration
  216. xprev, x = -1, guess
  217. while 1:
  218. t = x**(n - 1)
  219. xprev, x = x, ((n - 1)*x + y//t)//n
  220. if abs(x - xprev) < 2:
  221. break
  222. else:
  223. x = guess
  224. # Compensate
  225. t = x**n
  226. while t < y:
  227. x += 1
  228. t = x**n
  229. while t > y:
  230. x -= 1
  231. t = x**n
  232. return x, t == y
  233. def is_fermat_prp(n, a):
  234. if a < 2:
  235. raise ValueError("is_fermat_prp() requires 'a' greater than or equal to 2")
  236. if n < 1:
  237. raise ValueError("is_fermat_prp() requires 'n' be greater than 0")
  238. if n == 1:
  239. return False
  240. if n % 2 == 0:
  241. return n == 2
  242. a %= n
  243. if gcd(n, a) != 1:
  244. raise ValueError("is_fermat_prp() requires gcd(n,a) == 1")
  245. return pow(a, n - 1, n) == 1
  246. def is_euler_prp(n, a):
  247. if a < 2:
  248. raise ValueError("is_euler_prp() requires 'a' greater than or equal to 2")
  249. if n < 1:
  250. raise ValueError("is_euler_prp() requires 'n' be greater than 0")
  251. if n == 1:
  252. return False
  253. if n % 2 == 0:
  254. return n == 2
  255. a %= n
  256. if gcd(n, a) != 1:
  257. raise ValueError("is_euler_prp() requires gcd(n,a) == 1")
  258. return pow(a, n >> 1, n) == jacobi(a, n) % n
  259. def _is_strong_prp(n, a):
  260. s = bit_scan1(n - 1)
  261. a = pow(a, n >> s, n)
  262. if a == 1 or a == n - 1:
  263. return True
  264. for _ in range(s - 1):
  265. a = pow(a, 2, n)
  266. if a == n - 1:
  267. return True
  268. if a == 1:
  269. return False
  270. return False
  271. def is_strong_prp(n, a):
  272. if a < 2:
  273. raise ValueError("is_strong_prp() requires 'a' greater than or equal to 2")
  274. if n < 1:
  275. raise ValueError("is_strong_prp() requires 'n' be greater than 0")
  276. if n == 1:
  277. return False
  278. if n % 2 == 0:
  279. return n == 2
  280. a %= n
  281. if gcd(n, a) != 1:
  282. raise ValueError("is_strong_prp() requires gcd(n,a) == 1")
  283. return _is_strong_prp(n, a)
  284. def _lucas_sequence(n, P, Q, k):
  285. r"""Return the modular Lucas sequence (U_k, V_k, Q_k).
  286. Explanation
  287. ===========
  288. Given a Lucas sequence defined by P, Q, returns the kth values for
  289. U and V, along with Q^k, all modulo n. This is intended for use with
  290. possibly very large values of n and k, where the combinatorial functions
  291. would be completely unusable.
  292. .. math ::
  293. U_k = \begin{cases}
  294. 0 & \text{if } k = 0\\
  295. 1 & \text{if } k = 1\\
  296. PU_{k-1} - QU_{k-2} & \text{if } k > 1
  297. \end{cases}\\
  298. V_k = \begin{cases}
  299. 2 & \text{if } k = 0\\
  300. P & \text{if } k = 1\\
  301. PV_{k-1} - QV_{k-2} & \text{if } k > 1
  302. \end{cases}
  303. The modular Lucas sequences are used in numerous places in number theory,
  304. especially in the Lucas compositeness tests and the various n + 1 proofs.
  305. Parameters
  306. ==========
  307. n : int
  308. n is an odd number greater than or equal to 3
  309. P : int
  310. Q : int
  311. D determined by D = P**2 - 4*Q is non-zero
  312. k : int
  313. k is a nonnegative integer
  314. Returns
  315. =======
  316. U, V, Qk : (int, int, int)
  317. `(U_k \bmod{n}, V_k \bmod{n}, Q^k \bmod{n})`
  318. Examples
  319. ========
  320. >>> from sympy.external.ntheory import _lucas_sequence
  321. >>> N = 10**2000 + 4561
  322. >>> sol = U, V, Qk = _lucas_sequence(N, 3, 1, N//2); sol
  323. (0, 2, 1)
  324. References
  325. ==========
  326. .. [1] https://en.wikipedia.org/wiki/Lucas_sequence
  327. """
  328. if k == 0:
  329. return (0, 2, 1)
  330. D = P**2 - 4*Q
  331. U = 1
  332. V = P
  333. Qk = Q % n
  334. if Q == 1:
  335. # Optimization for extra strong tests.
  336. for b in bin(k)[3:]:
  337. U = (U*V) % n
  338. V = (V*V - 2) % n
  339. if b == "1":
  340. U, V = U*P + V, V*P + U*D
  341. if U & 1:
  342. U += n
  343. if V & 1:
  344. V += n
  345. U, V = U >> 1, V >> 1
  346. elif P == 1 and Q == -1:
  347. # Small optimization for 50% of Selfridge parameters.
  348. for b in bin(k)[3:]:
  349. U = (U*V) % n
  350. if Qk == 1:
  351. V = (V*V - 2) % n
  352. else:
  353. V = (V*V + 2) % n
  354. Qk = 1
  355. if b == "1":
  356. # new_U = (U + V) // 2
  357. # new_V = (5*U + V) // 2 = 2*U + new_U
  358. U, V = U + V, U << 1
  359. if U & 1:
  360. U += n
  361. U >>= 1
  362. V += U
  363. Qk = -1
  364. Qk %= n
  365. elif P == 1:
  366. for b in bin(k)[3:]:
  367. U = (U*V) % n
  368. V = (V*V - 2*Qk) % n
  369. Qk *= Qk
  370. if b == "1":
  371. # new_U = (U + V) // 2
  372. # new_V = new_U - 2*Q*U
  373. U, V = U + V, (Q*U) << 1
  374. if U & 1:
  375. U += n
  376. U >>= 1
  377. V = U - V
  378. Qk *= Q
  379. Qk %= n
  380. else:
  381. # The general case with any P and Q.
  382. for b in bin(k)[3:]:
  383. U = (U*V) % n
  384. V = (V*V - 2*Qk) % n
  385. Qk *= Qk
  386. if b == "1":
  387. U, V = U*P + V, V*P + U*D
  388. if U & 1:
  389. U += n
  390. if V & 1:
  391. V += n
  392. U, V = U >> 1, V >> 1
  393. Qk *= Q
  394. Qk %= n
  395. return (U % n, V % n, Qk)
  396. def is_fibonacci_prp(n, p, q):
  397. d = p**2 - 4*q
  398. if d == 0 or p <= 0 or q not in [1, -1]:
  399. raise ValueError("invalid values for p,q in is_fibonacci_prp()")
  400. if n < 1:
  401. raise ValueError("is_fibonacci_prp() requires 'n' be greater than 0")
  402. if n == 1:
  403. return False
  404. if n % 2 == 0:
  405. return n == 2
  406. return _lucas_sequence(n, p, q, n)[1] == p % n
  407. def is_lucas_prp(n, p, q):
  408. d = p**2 - 4*q
  409. if d == 0:
  410. raise ValueError("invalid values for p,q in is_lucas_prp()")
  411. if n < 1:
  412. raise ValueError("is_lucas_prp() requires 'n' be greater than 0")
  413. if n == 1:
  414. return False
  415. if n % 2 == 0:
  416. return n == 2
  417. if gcd(n, q*d) not in [1, n]:
  418. raise ValueError("is_lucas_prp() requires gcd(n,2*q*D) == 1")
  419. return _lucas_sequence(n, p, q, n - jacobi(d, n))[0] == 0
  420. def _is_selfridge_prp(n):
  421. """Lucas compositeness test with the Selfridge parameters for n.
  422. Explanation
  423. ===========
  424. The Lucas compositeness test checks whether n is a prime number.
  425. The test can be run with arbitrary parameters ``P`` and ``Q``, which also change the performance of the test.
  426. So, which parameters are most effective for running the Lucas compositeness test?
  427. As an algorithm for determining ``P`` and ``Q``, Selfridge proposed method A [1]_ page 1401
  428. (Since two methods were proposed, referred to simply as A and B in the paper,
  429. we will refer to one of them as "method A").
  430. method A fixes ``P = 1``. Then, ``D`` defined by ``D = P**2 - 4Q`` is varied from 5, -7, 9, -11, 13, and so on,
  431. with the first ``D`` being ``jacobi(D, n) == -1``. Once ``D`` is determined,
  432. ``Q`` is determined to be ``(P**2 - D)//4``.
  433. References
  434. ==========
  435. .. [1] Robert Baillie, Samuel S. Wagstaff, Lucas Pseudoprimes,
  436. Math. Comp. Vol 35, Number 152 (1980), pp. 1391-1417,
  437. https://doi.org/10.1090%2FS0025-5718-1980-0583518-6
  438. http://mpqs.free.fr/LucasPseudoprimes.pdf
  439. """
  440. for D in range(5, 1_000_000, 2):
  441. if D & 2: # if D % 4 == 3
  442. D = -D
  443. j = jacobi(D, n)
  444. if j == -1:
  445. return _lucas_sequence(n, 1, (1-D) // 4, n + 1)[0] == 0
  446. if j == 0 and D % n:
  447. return False
  448. # When j == -1 is hard to find, suspect a square number
  449. if D == 13 and is_square(n):
  450. return False
  451. raise ValueError("appropriate value for D cannot be found in is_selfridge_prp()")
  452. def is_selfridge_prp(n):
  453. if n < 1:
  454. raise ValueError("is_selfridge_prp() requires 'n' be greater than 0")
  455. if n == 1:
  456. return False
  457. if n % 2 == 0:
  458. return n == 2
  459. return _is_selfridge_prp(n)
  460. def is_strong_lucas_prp(n, p, q):
  461. D = p**2 - 4*q
  462. if D == 0:
  463. raise ValueError("invalid values for p,q in is_strong_lucas_prp()")
  464. if n < 1:
  465. raise ValueError("is_selfridge_prp() requires 'n' be greater than 0")
  466. if n == 1:
  467. return False
  468. if n % 2 == 0:
  469. return n == 2
  470. if gcd(n, q*D) not in [1, n]:
  471. raise ValueError("is_strong_lucas_prp() requires gcd(n,2*q*D) == 1")
  472. j = jacobi(D, n)
  473. s = bit_scan1(n - j)
  474. U, V, Qk = _lucas_sequence(n, p, q, (n - j) >> s)
  475. if U == 0 or V == 0:
  476. return True
  477. for _ in range(s - 1):
  478. V = (V*V - 2*Qk) % n
  479. if V == 0:
  480. return True
  481. Qk = pow(Qk, 2, n)
  482. return False
  483. def _is_strong_selfridge_prp(n):
  484. for D in range(5, 1_000_000, 2):
  485. if D & 2: # if D % 4 == 3
  486. D = -D
  487. j = jacobi(D, n)
  488. if j == -1:
  489. s = bit_scan1(n + 1)
  490. U, V, Qk = _lucas_sequence(n, 1, (1-D) // 4, (n + 1) >> s)
  491. if U == 0 or V == 0:
  492. return True
  493. for _ in range(s - 1):
  494. V = (V*V - 2*Qk) % n
  495. if V == 0:
  496. return True
  497. Qk = pow(Qk, 2, n)
  498. return False
  499. if j == 0 and D % n:
  500. return False
  501. # When j == -1 is hard to find, suspect a square number
  502. if D == 13 and is_square(n):
  503. return False
  504. raise ValueError("appropriate value for D cannot be found in is_strong_selfridge_prp()")
  505. def is_strong_selfridge_prp(n):
  506. if n < 1:
  507. raise ValueError("is_strong_selfridge_prp() requires 'n' be greater than 0")
  508. if n == 1:
  509. return False
  510. if n % 2 == 0:
  511. return n == 2
  512. return _is_strong_selfridge_prp(n)
  513. def is_bpsw_prp(n):
  514. if n < 1:
  515. raise ValueError("is_bpsw_prp() requires 'n' be greater than 0")
  516. if n == 1:
  517. return False
  518. if n % 2 == 0:
  519. return n == 2
  520. return _is_strong_prp(n, 2) and _is_selfridge_prp(n)
  521. def is_strong_bpsw_prp(n):
  522. if n < 1:
  523. raise ValueError("is_strong_bpsw_prp() requires 'n' be greater than 0")
  524. if n == 1:
  525. return False
  526. if n % 2 == 0:
  527. return n == 2
  528. return _is_strong_prp(n, 2) and _is_strong_selfridge_prp(n)