Solution to LeetCode Problem 204: Count Primes

count_primes.py

"""Utilities for Counting and Testing Prime Numbers.

A significant portion of the material here is thanks to Al Sweigart and their
incredible explanations on prime number sieves from their book, "Cracking Codes
With Python."

You can read the book online here:
  * https://inventwithpython.com/cracking/chapter0.html

You can also purchase a paper copy here:
  * https://www.amazon.com/Cracking-Codes-Python-Introduction-Building/dp/1593278225

The specific material on prime number sieves is located in Chapter 22:
  * https://inventwithpython.com/cracking/chapter22.html
"""

from enum import Enum, auto, unique
from itertools import count, takewhile
from math import isqrt
from typing import Callable, Iterable, List

import itertools
import unittest
import random


def dummyPrimeCountingMethod(n: int) -> int:
  return -1


def candidatePrimeDivisorsFor(n: int) -> Iterable[int]:
  for candidate_divisor in takewhile(lambda x: x <= isqrt(n), count(3, 2)):
    yield candidate_divisor


def nIsDivisibleBy(n: int, m: int) -> bool:
  return n % m == 0


def isEven(n: int) -> bool:
  return nIsDivisibleBy(n, 2)


def isPrimeBasedOnTrialDivision(n: int) -> bool:
  if n < 2:
    return False

  if n in {2, 3, 5, 7}:
    return True

  if isEven(n):
    return False

  for i in candidatePrimeDivisorsFor(n):
    if nIsDivisibleBy(n, i):
      return False

  return True


def countPrimesUsingTrialDivision(n: int) -> int:
  return sum(isPrimeBasedOnTrialDivision(i) for i in range(n))


def countPrimesUsingSieveOfEratosthenes(n: int) -> int:
  # Zero and one do not fit into the definition of prime numbers because
  # nothing is divisible by zero and zero divided by anything is zero, and we
  # exclude the number one for definitional purposes.
  #
  # In other words, one could technically be a prime number, but including it
  # in the definition would make a ton of other definitions significantly more
  # complicated without actually adding any value or usefulness, so we just
  # leave it out.
  sieve = [False, False] + [True for _ in range(2, n)]

  for index, value in enumerate(sieve):
    # There's nothing for us to do here if we've already excluded this number
    # as a potential prime during a previous iteration.
    if value is False:
      continue

    # Since this number hasn't been marked as composite, we now know this
    # number is prime. We now need to mark off all of its multiples in the list
    # all the way until the end of the list.
    for m_index in range(2 * index, n, index):
      sieve[m_index] = False

  # Every number that didn't get marked in the for loop is a confirmed prime,
  # so we just need to return the count of how many numbers we ended up with.
  return sum(sieve)


def EratosthenesPrimeNumberSieve(n: int) -> Iterable[int]:
  sieve = [True] * n

  if n >= 1:
    sieve[0] = False
  if n >= 2:
    sieve[1] = False

  # Create the prime number sieve.
  for i in range(2, isqrt(n) + 1):
    pointer = 2 * i

    while pointer < n:
      sieve[pointer] = False
      pointer += i

  # Compile list of all of the prime numbers strictly less than n.
  for i in range(n):
    if sieve[i] == True:
      yield i


def millerRabinPrimalityTest(n: int, iterations: int = 20) -> bool:
  """Returns true if n is a prime number.

  The Miller-Rabin primality test is a probabilistic primality test. The test
  output indicates that either a number is definitely composite or "probably
  prime."

  The probability of falsely categorizing a composite number as prime is at
  most equal to 4^(-k), where k is the number of iterations of the test
  performed.

  Additional Reading:
    * https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test
  """
  if n < 2 or isEven(n):
    return False

  if n == 3:
    return True

  s = n - 1
  t = 0

  while isEven(s):
    # Keep halving s until s is odd, using t to keep track of how many times we
    # have to perform this operation.
    s //= 2
    t += 1

  # Attempt to falsify the hypothesis that n is prime `iterations` number of
  # times.
  for trial in range(iterations):
    a = random.randrange(2, n - 1)
    v = pow(a, s, n)

    # This test does not apply if `v` is equal to one.
    if v != 1:
      i = 0

      while v != n - 1:
        if i == t - 1:
          return False

        i += 1
        v = (v ** 2) % n

  return True


def isPrime(n: int, iterations: int = 20, size: int = 100) -> bool:
  if n < 2:
    return False

  for prime in EratosthenesPrimeNumberSieve(size):
    if n == prime:
      return True
    if nIsDivisibleBy(n, prime):
      return False

  # If the above checks weren't able to figure out whether the number is prime
  # or not, we resort to the Miller-Rabin test.
  return millerRabinPrimalityTest(iterations)


def generatePrime(bits: int = 64) -> int:
  """Returns randomly-generated prime number of the given size."""
  while True:
    n = random.randrange(2**(bits - 1), 2**(bits))

    # The number of Miller-Rabin iterations needs to be pretty much no higher
    # than 5-7 in order for this function to terminate before you've visibily
    # aged.
    if isPrime(n, 7):
      return n


def generatePrimes() -> Iterable[int]:
  """Generate an infinite sequence of prime numbers.

  The credit for this infinite prime number generator goes to Eli Bendersky,
  who's personal website I've linked below.

  Source:
   * https://eli.thegreenplace.net/2023/my-favorite-prime-number-generator/
  """
  D = {}
  q = 2

  while True:
    if q not in D:
      D[q * q] = [q]
      yield q
    else:
      for p in D[q]:
        D.setdefault(p + q, []).append(p)
      del D[q]
    q += 1


def generatePrimesOptimized() -> Iterable[int]:
  """Generates an infinite sequence of prime numbers.

  I also got this version from Eli Bendersky's website, and according to him,
  this version of the infinite primes generator was written by Alex Martelli,
  Tim Hochberg, and Wolfgang Beneicke in an ActiveState Recipe thread.

  Source:
   * https://eli.thegreenplace.net/2023/my-favorite-prime-number-generator/
  """
  yield 2
  D = {}
  for q in itertools.count(3, step=2):
    p = D.pop(q, None)
    if not p:
      D[q * q] = q
      yield q
    else:
      x = q + p + p  # get odd multiples
      while x in D:
        x += p + p
      D[x] = p


@unique
class PrimeCountingMethod(Enum):
  DUMMY_METHOD = auto()
  TRIAL_DIVISION = auto()
  SIEVE_OF_ERATOSTHENES = auto()


def getPrimeCountingMethod(method: PrimeCountingMethod) -> Callable[[int], int]:
  match method:
    case PrimeCountingMethod.TRIAL_DIVISION:
      return countPrimesUsingTrialDivision
    case PrimeCountingMethod.SIEVE_OF_ERATOSTHENES:
      return countPrimesUsingSieveOfEratosthenes
    case _:
      return dummyPrimeCountingMethod


def getDefaultPrimeCountingMethod() -> PrimeCountingMethod:
  return PrimeCountingMethod.SIEVE_OF_ERATOSTHENES


def numberOfPrimesLessThan(n: int, method: PrimeCountingMethod = getDefaultPrimeCountingMethod()) -> int:
  prime_counting_function = getPrimeCountingMethod(method)
  return prime_counting_function(n)


class Solution:
  def countPrimes(self, n: int) -> int:
    return numberOfPrimesLessThan(n)


class TestSolution(unittest.TestCase):
  def testPrimeSieve(self) -> None:
    with self.subTest('10'):
      primes = [p for p in EratosthenesPrimeNumberSieve(10)]
      print(primes)
      number_of_primes = len(primes)
      expected_number_of_primes = 4
      self.assertEqual(number_of_primes, expected_number_of_primes)

    with self.subTest('0'):
      primes = [p for p in EratosthenesPrimeNumberSieve(0)]
      print(primes)
      number_of_primes = len(primes)
      expected_number_of_primes = 0
      self.assertEqual(number_of_primes, expected_number_of_primes)

    with self.subTest('150000'):
      primes = [p for p in EratosthenesPrimeNumberSieve(150000)]
      print(primes)
      number_of_primes = len(primes)
      expected_number_of_primes = 13848
      self.assertEqual(number_of_primes, expected_number_of_primes)


if __name__ == "__main__":
  unittest.main()