feat(math.py): Add the functionality to generate a random prime number within a specified range

cyaron/math.py

@@ -25,6 +25,9 @@
     miu(x): The MIU function of x
     dec2base(n,base): Number base conversion
     n2words(num,join=True): Number to words
+    nextprime(n): Find the next prime number after n
+    prevprime(n,raise_error=True): Find the previous prime number before n
+    randprime(a,b,raise_error=True): Generate a random prime number in the range [a,b]
 
 forked from https://blog.dreamshire.com/common-functions-routines-project-euler/
 """
@@ -549,3 +552,94 @@ def n2words(num: int, join: bool = True):
     if join:
         return ' '.join(words)
     return words
+
+
+def nextprime(n: int):
+    """
+    Find the next prime number after a given number.
+    Args:
+        n: The number after which to find the next prime.
+    Returns:
+        The next prime number after n.
+    """
+
+    if n < 2:
+        return 2
+    if n == 2:
+        return 3
+    if n % 6 == 0:
+        if miller_rabin(n + 1):
+            return n + 1
+        n += 6
+    elif 1 <= n % 6 < 5:
+        n += 6 - n % 6
+    else:
+        if miller_rabin(n + 2):
+            return n + 2
+        n += 7
+    while True:
+        if miller_rabin(n - 1):
+            return n - 1
+        if miller_rabin(n + 1):
+            return n + 1
+        n += 6
+
+
+def prevprime(n: int, raise_error: bool = True):
+    """
+    Find the previous prime number before a given number.
+    Args:
+        n: The number before which to find the previous prime.
+        raise_error: Raise an error if there is no prime number less than `n`. Default is True.
+    Returns:
+        The previous prime number before `n`.
+    Raises:
+        ValueError: If there is no prime number less than `n` and `raise_error` is True.
+    """
+    if n <= 2:
+        if raise_error:
+            raise ValueError(f"No prime number less than {n}")
+        return 0
+    if n == 3:
+        return 2
+    if n <= 5:
+        return 3
+    if n % 6 == 0:
+        if miller_rabin(n - 1):
+            return n - 1
+        n -= 6
+    elif n % 6 == 1:
+        if miller_rabin(n - 2):
+            return n - 2
+        n -= 7
+    else:
+        n -= n % 6
+    while True:
+        if miller_rabin(n + 1):
+            return n + 1
+        if miller_rabin(n - 1):
+            return n - 1
+        n -= 6
+
+
+def randprime(a: int, b: int, raise_error: bool = True):
+    """
+    Generate a random prime number in the range [a, b].
+    Args:
+        a: The lower bound of the range.
+        b: The upper bound of the range.
+    Returns:
+        A random prime number in the specified range.
+    """
+    st = random.randint(a, b)
+    if miller_rabin(st):
+        return st
+    nxt = nextprime(st)
+    if nxt <= b:
+        return nxt
+    pre = prevprime(st)
+    if pre >= a:
+        return pre
+    if raise_error:
+        raise ValueError(f"No prime number in the range [{a}, {b}]")
+    return 0

cyaron/tests/__init__.py

@@ -6,3 +6,4 @@
 from .graph_test import TestGraph
 from .vector_test import TestVector
 from .general_test import TestGeneral
+from .math_test import TestMath

cyaron/tests/math_test.py

@@ -0,0 +1,46 @@
+import unittest
+from random import randint
+from cyaron import randprime, nextprime, prevprime, prime_sieve
+
+
+class TestMath(unittest.TestCase):
+
+    def setUp(self):
+        self.prime_range = 100000
+        self.prime_set = set(prime_sieve(self.prime_range * 2))
+        return super().setUp()
+
+    def test_randprime(self):
+        for _ in range(20):
+            self.assertTrue(randprime(1, self.prime_range) in self.prime_set)
+
+    def test_nextprime(self):
+
+        def bf_nextprime(n):
+            while True:
+                n += 1
+                if n in self.prime_set:
+                    return n
+
+        for i in range(20):
+            n = randint(1, self.prime_range)
+            self.assertEqual(nextprime(n), bf_nextprime(n))
+            self.assertEqual(nextprime(i), bf_nextprime(i))
+
+    def test_prevprime(self):
+
+        def bf_prevprime(n):
+            while True:
+                n -= 1
+                if n in self.prime_set:
+                    return n
+                if n < 2:
+                    return 0
+
+        for i in range(20):
+            n = randint(1, self.prime_range)
+            self.assertEqual(prevprime(n, False), bf_prevprime(n))
+            self.assertEqual(prevprime(i, False), bf_prevprime(i))
+
+        self.assertRaises(ValueError, prevprime, 1)
+        self.assertRaises(ValueError, prevprime, 2)