Skip to content

Commit f437031

Browse files
authored
Faster sieve() recipe (python#98287)
1 parent 146f168 commit f437031

File tree

1 file changed

+27
-8
lines changed

1 file changed

+27
-8
lines changed

Doc/library/itertools.rst

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -809,15 +809,25 @@ which incur interpreter overhead.
809809
for k in range(len(roots) + 1)
810810
]
811811

812+
def iter_index(seq, value, start=0):
813+
"Return indices where a value occurs in a sequence."
814+
# iter_index('AABCADEAF', 'A') --> 0 1 4 7
815+
i = start - 1
816+
try:
817+
while True:
818+
yield (i := seq.index(value, i+1))
819+
except ValueError:
820+
pass
821+
812822
def sieve(n):
813-
"Primes less than n"
814-
# sieve(30) --> 2 3 5 7 11 13 17 19 23 29
815-
data = bytearray([1]) * n
816-
data[:2] = 0, 0
817-
limit = math.isqrt(n) + 1
818-
for p in compress(range(limit), data):
819-
data[p+p : n : p] = bytearray(len(range(p+p, n, p)))
820-
return compress(count(), data)
823+
"Primes less than n"
824+
# sieve(30) --> 2 3 5 7 11 13 17 19 23 29
825+
data = bytearray([1]) * n
826+
data[:2] = 0, 0
827+
limit = math.isqrt(n) + 1
828+
for p in compress(range(limit), data):
829+
data[p*p : n : p] = bytearray(len(range(p*p, n, p)))
830+
return iter_index(data, 1)
821831

822832
def flatten(list_of_lists):
823833
"Flatten one level of nesting"
@@ -1170,6 +1180,15 @@ which incur interpreter overhead.
11701180
>>> all(factored(x) == expanded(x) for x in range(-10, 11))
11711181
True
11721182

1183+
>>> list(iter_index('AABCADEAF', 'A'))
1184+
[0, 1, 4, 7]
1185+
>>> list(iter_index('AABCADEAF', 'B'))
1186+
[2]
1187+
>>> list(iter_index('AABCADEAF', 'X'))
1188+
[]
1189+
>>> list(iter_index('', 'X'))
1190+
[]
1191+
11731192
>>> list(sieve(30))
11741193
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
11751194
>>> small_primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

0 commit comments

Comments
 (0)