Count Primes Python Solution with Tests

Tested Python solution for LeetCode 204 with 18 pytest cases. Generate a practice environment with lcpy.

LeetCode 204, Medium. Topics: Array, Math, Enumeration, Number Theory, Primality Test, Sieve Theory, Prime Number Sieve. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 18 parametrized pytest cases, and a playground notebook:

lcpy gen -n 204   # by problem number
lcpy gen -s count_primes   # by problem name

​

Problem

Given an integer n, return the number of prime numbers that are strictly less than n.

​

Examples

Input: n = 10
Output: 4
Explanation: There are 4 prime numbers less than 10, they are 2, 3, 5, 7.

Input: n = 0
Output: 0

Input: n = 1
Output: 0

​

Constraints

• 0 <= n <= 5 * 10⁶

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

class Solution:
    # Time: O(n log log n)
    # Space: O(n)
    def count_primes(self, n: int) -> int:
        if n < 3:
            return 0
        is_prime = [True] * n
        is_prime[0] = is_prime[1] = False
        for i in range(2, int(n**0.5) + 1):
            if is_prime[i]:
                for multiple in range(i * i, n, i):
                    is_prime[multiple] = False
        return sum(is_prime)

​

Complexity

Time	Space
O(n log log n)	O(n)

​

Tags

NeetCode All.