Artemis: Optimize Prime Number Functions and Improve Documentation

src/llm_benchmark/algorithms/primes.py

@@ -1,10 +1,12 @@
+import typing
 from typing import List
 
 
 class Primes:
     @staticmethod
     def is_prime(n: int) -> bool:
-        """Check if a number is prime
+        """
+        Check if a number is prime
 
         Args:
             n (int): Number to check
@@ -14,14 +16,15 @@ def is_prime(n: int) -> bool:
         """
         if n < 2:
             return False
-        for i in range(2, n):
+        for i in range(2, int(n ** 0.5) + 1):
             if n % i == 0:
                 return False
         return True
 
     @staticmethod
     def sum_primes(n: int) -> int:
-        """Sum of primes from 0 to n (exclusive)
+        """
+        Sum of primes from 0 to n (exclusive)
 
         Args:
             n (int): Number to sum up to
@@ -30,14 +33,15 @@ def sum_primes(n: int) -> int:
             int: Sum of primes from 0 to n
         """
         sum_ = 0
-        for i in range(n):
+        for i in range(2, n):
             if Primes.is_prime(i):
                 sum_ += i
         return sum_
 
     @staticmethod
     def prime_factors(n: int) -> List[int]:
-        """Prime factors of a number
+        """
+        Prime factors of a number
 
         Args:
             n (int): Number to factorize
@@ -47,9 +51,12 @@ def prime_factors(n: int) -> List[int]:
         """
         ret = []
         while n > 1:
-            for i in range(2, n + 1):
+            for i in range(2, int(n ** 0.5) + 1):
                 if n % i == 0:
                     ret.append(i)
                     n = n // i
                     break
-        return ret
+            else:
+                ret.append(n)
+                n = 1
+        return ret