forked from sympy/sympy
/
generate.py
172 lines (143 loc) · 4.47 KB
/
generate.py
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
169
170
171
172
"""
Generating and counting primes.
"""
import random
from primetest import isprime
# Using arrays for sieving instead of lists greatly reduces
# memory consumption
from array import array as _array
def _arange(a, b):
ar = _array('l', [0]*(b-a))
for i, e in enumerate(xrange(a, b)):
ar[i] = e
return ar
class Sieve:
"""An infinite list of prime numbers, implemented as a dynamically
growing sieve of Eratosthenes. When a lookup is requested involving
a number that has not been sieved, the sieve is automatically
extended up to that number."""
_list = _array('l', [2, 3, 5, 7, 11, 13])
def __repr__(self):
return "<Sieve with %i primes sieved: 2, 3, 5, ... %i, %i>" % \
(len(self._list), self._list[-2], self._list[-1])
def extend(self, N):
"""Grow the sieve to cover all numbers <= N."""
if N <= self._list[-1]:
return
# We need to sieve against all bases up to sqrt(n). If there
# are too few, extend the list recursively.
maxbase = int(N**0.5)+1
self.extend(maxbase)
# Create a new sieve starting from N**0.5
begin = self._list[-1] + 1
newsieve = _arange(begin, N+1)
# Now eliminate all multiples of primes in [2, N**0.5]
for p in self.primerange(2, maxbase):
# Start counting at a multiple of p, offsetting
# the index to account for the new sieve's base index
startindex = (-begin) % p
for i in xrange(startindex, len(newsieve), p):
newsieve[i] = 0
# Merge the sieves
self._list += _array('l', [x for x in newsieve if x])
def extend_to_no(self, n):
"""Extend to include (at least) the nth prime number"""
while len(self._list) < n:
self.extend(int(self._list[-1] * 1.5))
def primerange(self, a, b):
"""Generate all prime numbers in the range [a, b)."""
assert a <= b
if b < 2:
return
a = max(2, a)
self.extend(b)
i = self.search(a)[1]
maxi = len(self._list) + 1
while i < maxi:
p = self._list[i-1]
if p < b:
yield p
i += 1
else:
return
def search(self, n):
"""For n >= 2, return the tightest a, b such that
self[a] <= n <= self[b]"""
assert n >= 2
if n > self._list[-1]:
self.extend(n)
# Use a binary search
a = 0
b = len(self._list)-1
while 1:
m = (a + b) // 2
if self._list[m] == n:
return m+1, m+1
elif self._list[m] > n:
b = m - 1
elif self._list[m] < n:
a = m + 1
if a > b:
return b+1, a+1
def __contains__(self, n):
if n < 2:
return False
a, b = self.search(n)
return a == b
def __getitem__(self, n):
"""Return the nth prime number"""
self.extend_to_no(n)
return self._list[n-1]
# Generate a global object for repeated use in trial division etc
sieve = Sieve()
def prime(n):
"""Return the nth prime, with the primes indexed as prime(1) = 2,
prime(2) = 3, etc."""
assert n > 0
return sieve[n]
def primepi(n):
"""
Return the value of the prime counting function pi(n) = the number
of prime numbers less than or equal to n. The number n need not
necessarily be an integer.
"""
if n < 2:
return 0
else:
n = int(n)
return sieve.search(n)[0]
# TODO: the following functions could use the sieve for small n
def nextprime(n):
"""Return the smallest prime greater than n."""
n = max(n, 0)
while 1:
n += 1
if isprime(n):
return n
def prevprime(n):
"""Return the largest prime smaller than n."""
if n < 3:
raise ValueError("no preceding primes")
while 1:
n -= 1
if isprime(n):
return n
def primerange(a, b):
"""Generate a list of all prime numbers in the range [a, b)."""
assert a <= b
a -= 1
while 1:
a = nextprime(a)
if a < b:
yield a
else:
return
def randprime(a, b):
"""Return a random prime number in the range [a, b)"""
n = random.randint(a-1, b)
p = nextprime(n)
if p >= b:
p = prevprime(b)
if p < a:
raise ValueError("no primes exist in the specified range")
return p