math/primes.py at master · NepNepFFXIV/math · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def time_elapsed(func):
    import time
    from functools import wraps
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        print(f"{func.__name__} finished in {(end_time - start_time):.5f}s")
        return result
    return wrapper

# Trial division using 6k+-1 optimization
def is_prime(n):
    if n <= 3: return n > 1
    if n % 2 == 0 or n % 3 == 0: return False

    for i in range(5, int(n**0.5)+1, 6):
        if n % i == 0 or n % (i+2) == 0:
            return False
    return True

# Brute force trial division
def is_prime_brute(n):
    if n < 2: return False
    for i in range(2, int(n**0.5)+1):
        if n % i == 0:
            return False
    return True

# prime_checker is used to check if a number is prime
def prime_generator(prime_checker=is_prime):
    yield 2
    n = 1
    while True:
        n += 2
        if prime_checker(n):
            yield(n)

# See sieve_eratosthenes_2 implementation
def prime_generator_eratosthenes(n=15_000_000): # arbitrary limit
    marked = [False, True] * (n+2>>1)
    marked[2] = True
    for p in range(3, int(n**0.5)+1, 2):
        if marked[p]:
            for i in range(p*p, n+1, 2*p):
                marked[i] = False
    yield 2
    for p in range(3,n+1, 2):
        if marked[p]:
            yield p

# All primes less than or equal n - O(n log n)
@time_elapsed
def sieve_sundaram(n):
    primes = []
    m = n//2
    marked = [False] * (m + 1)

    # Main logic of Sundaram. Mark all
    # numbers which do not generate prime
    for i in range(1, m + 1):
        for j in range((i * (i+1)) << 1, m + 1, 2*i + 1):
            marked[j] = True

    primes.append(2)

    for i in range(1, m):
        if (marked[i] == False):
            primes.append(2*i + 1)

    if is_prime(n): primes.append(n)
    return primes

# All primes less than or equal n - O(n log log n) -
# Eratosthenes - No Optimization
# https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes
@time_elapsed
def sieve_eratosthenes_1(n):
    marked = [True] * (n+1)
    for p in range(2, int(n**0.5)+1):
        if marked[p]:                       # marked[p] == True
            for i in range(p*p, n+1, p):    # Update all multiplies of p
                marked[i] = False
    return [p for p in range(2,n+1) if marked[p]]

# Eratosthenes - Odds optimization
@time_elapsed
def sieve_eratosthenes_2(n):
    marked = [False, True] * (n+2>>1)
    marked[2] = True
    for p in range(3, int(n**0.5)+1, 2):
        if marked[p]:                         # marked[p] == True
            for i in range(p*p, n+1, 2*p):    # Update all multiplies of p. No need to check for multiple 2 of p hence increase by 2*p each step
                marked[i] = False
    return [2] + [p for p in range(3, n+1, 2) if marked[p]]

# Eratosthenes - Odds optimization with help of Numpy masking to save memmory
@time_elapsed
def sieve_eratosthenes_3(n):
    import numpy as np
    marked = np.array([False, True] * (n+2>>1))
    marked[2] = True
    for p in range(3, int(n**0.5)+1, 2):
        if marked[p]:
            marked[p*p : n+1: 2*p] = False
    return [2] + [p for p in range(3, n+1, 2) if marked[p]]

# All primes less than or equal n - O(n / (log log n))
@time_elapsed
def sieve_atkin(n):
    #idk lol
    pass

# All primes less than or equal n using prime generator - O(lol)
@time_elapsed
def primes_lte_n(n, generator=prime_generator_eratosthenes):
    nums = []
    gen = generator()
    m = next(gen)
    while m <= n:
        nums.append(m)
        m = next(gen)
    return nums

# Prime Factorization - Optimized Trial division - O(sqrt(n))
def prime_factorization(n):
    result = []
    while n & 1 == 0:
        result.append(2)
        n = n // 2
    for i in range(3, int(n**0.5)+1, 2):
        while n % i == 0:
            result.append(i)
            n = n // i
    if n > 2: result.append(n)
    return result

# All possible pairs of Goldbach's conjecture
def goldbach(n, sieve=sieve_eratosthenes_2):
    # We only check even numbers greater than 2
    if n < 2 or n & 1 == 1: return []

    primes = sieve(n)
    result = []
    for i in primes:
        difference = n - i
        if difference < i: break
        if difference in primes:
            result.append([i, difference])
    return result

if __name__ == "__main__":
    # n = int(math.pow(10, 7))
    n = 20_000_000
    # print(is_prime(n))
    # print(is_prime_brute(n))
    # print(prime_factorization(n))
    # print(primes_lte_n(n))
    # print(sieve_sundaram(n))
    # sieve_atkin(n)
    # print(goldbach(n))
    # print(is_prime(n))
    # print(is_prime_wheel_factorization(n))

    sieve_eratosthenes_1(n)
    sieve_eratosthenes_2(n)
    sieve_eratosthenes_3(n)