TL;DR: Just use PyPy; it gets you to about 10x the time of C++. If you really want to use CPython, a lot of clever optimizations (not algorithm changes) gets you as fast as PyPy and then using Numpy gets you close to C++ (2x the time).
The first thing of note is that your Python code is broken:
if m > len(marked):
break
Remember that Python is 0-indexed. What about when m == len(marked)?
So that's the first thing to fix. Going from the top, I'd do the
translation so:
from math import log
n = 10001 n <- 10001
maximum = n * (log(n) + log(log(n))) max <- n*(log(n) + log(log(n)))
marked = [False] * int(maximum // 2) marked <- vector(mode="logical", length=max/2)
for i in range(1, len(marked)+1): for (i in 1:length(marked)) {
for j in range(1, i+1): for (j in 1:i) {
m = i + j + 2*i*j m <- i + j + 2*i*j
if m > len(marked): break if (m > length(marked)) break
marked[m-1] = True marked[m] <- T
}
}
count = 1 count <- 1
for i in range(1, len(marked)+1): for (i in 1:length(marked)) {
if not marked[i-1]: count += 1 if (!marked[i]) count <- count + 1
if count == n: if (count == n) {
print(2*i + 1) print(2*i + 1)
break break
}
}
This is as direct a mapping as possible; instead of changing the comparison
I just shifted the index when indexing. This isn't idiomatic, but it's direct.
It's largely the same as your code, but it's correct. This matters when we go to
larger N, where your code fails. It's also Python 3 compatible simply by using
print with brackets.
Let's time them:
$ time Rscript r.r
[1] 104743
Rscript r.r 1.59s user 0.78s system 99% cpu 2.375 total
$ time python2 p.py
104743
python2 p.py 12.88s user 0.00s system 100% cpu 12.873 total
$ time python3 p.py
104743
python3 p.py 0.16s user 0.00s system 98% cpu 0.163 total
$ # A faster, very compatible Python interpreter
$ time pypy p.py
104743
pypy p.py 0.04s user 0.01s system 98% cpu 0.051 total
$ time pypy3 p.py
104743
pypy3 p.py 0.05s user 0.01s system 99% cpu 0.054 total
So there we have it. Python is over an order of magnitude faster on 75% of
interpreters, and under an order of magnitude slower in the worst case...
But why is it so slow with python2? line_profiler is a good utility:
import line_profiler
profiler = line_profiler.LineProfiler()
def main():
# ... the code ...
profiler.enable()
profiler.add_function(main)
main()
profiler.print_stats()
Giving:
Line # Hits Time Per Hit % Time Line Contents
==============================================================
4 def main():
5 1 271 271.0 0.0 from math import log
6
7 1 1 1.0 0.0 n = 10001
8 1 22 22.0 0.0 maximum = n * (log(n) + log(log(n)))
9 1 373 373.0 0.0 marked = [False] * int(maximum // 2)
10 57160 37122 0.6 0.3 for i in range(1, len(marked)+1):
11 188523 8444200 44.8 62.9 for j in range(1, i+1):
12 188355 118260 0.6 0.9 m = i + j + 2*i*j
13 188355 4672922 24.8 34.8 if m > len(marked): break
14 131364 75962 0.6 0.6 marked[m-1] = True
15
16
17 1 1 1.0 0.0 count = 1
18 52371 25618 0.5 0.2 for i in range(1, len(marked)+1):
19 52371 28393 0.5 0.2 if not marked[i-1]: count += 1
20 52371 26003 0.5 0.2 if count == n:
21 1 60 60.0 0.0 print(2*i + 1)
22 1 329 329.0 0.0 break
So our most likely offending line is:
11 188523 8444200 44.8 62.9 for j in range(1, i+1):
On Python 2 there is both range and xrange. Using
for i in range(...):
will generate a list of numbers and then loop over them, whereas
for i in xrange(...):
will just increment i as wanted each time. This is often much faster. So let's add:
# Set range to xrange on Python 2
try:
range = xrange
except NameError:
pass
to the top of the script and retime:
$ time python2 p.py
104743
python2 p.py 0.09s user 0.00s system 96% cpu 0.090 total
$ time python3 p.py
104743
python3 p.py 0.16s user 0.00s system 99% cpu 0.161 total
$ time pypy p.py
104743
pypy p.py 0.04s user 0.00s system 93% cpu 0.050 total
$ time pypy3 p.py
104743
pypy3 p.py 0.04s user 0.01s system 99% cpu 0.054 total
Yup. CPython is much improved.
So how would we make the code good? First you can transform the i+j+2ij
calculation to directly use a stepping range:
step = 2*i + 1
for m in range(i+step-1, i+i*step, step):
if m >= len(marked):
break
marked[m] = True
Then you can just do a slice assignment:
for i in range(1, len(marked)+1):
step = 2*i + 1
marked[i+step-1::step] = [True] * len(marked[i+step-1::step])
and you can immediately see that you can stop when...
if i+step-1 >= len(marked):
break
This leads to calculating the bound in the for loop itself.
step = 2*i + 1, so i+step-1 = i+(2*i+1)-1 = 3*i, so we should
stop before len(marked)//3:
for i in range(1, len(marked)//3):
step = 2*i + 1
marked[i+step-1::step] = [True] * len(marked[i+step-1::step])
The count part can use enumerate:
count = 1
for i, mark in enumerate(marked, 1):
if not mark:
count += 1
if count == n:
print(2*i + 1)
break
I suggest inverting the True/False in the array to make this simpler,
and then you can just do count += mark. Why? Well, it lets us use
itertools.accumulate on Python 3:
for i, count in enumerate(accumulate(marked), 1):
if count+1 == n:
print(2*i + 1)
break
You might also think len(marked[i+step-1::step]) is wasteful, since
it's generating a new array each time: it is!
It's still really fast, but Jaime's solution is better... except it
involves risky math steps. How could we get the best of both?
Well, remember how in Python 3 range is a fancy "pseudo-list" and
doesn't build a real list? Python 3's range can be sliced, too.
We just need to build a range the same length as the list and slice
that instead. This is actually slower on CPython (the one invoked
as python3) for normal-sized lists because CPython's list slices
are really fast, but on PyPy the overheads of the range version
are optimized out:
from itertools import accumulate
from math import log
n = 1000001
maximum = n * (log(n) + log(log(n)))
marked = [True] * int(maximum // 2)
slice_proxy = range(len(marked))
for i in range(1, len(marked)//3):
step = 2*i + 1
marked[i+step-1::step] = [False] * len(slice_proxy[i+step-1::step])
for i, count in enumerate(accumulate(marked), 1):
if count+1 == n:
print(2*i + 1)
break
(Note the much larger n)
Times without slice_proxy:
$ time python3 p.py
15485867
python3 p.py 5.57s user 0.01s system 100% cpu 5.576 total
$ time pypy3 p.py
15485867
pypy3 p.py 3.58s user 0.08s system 100% cpu 3.654 total
Times with slice_proxy:
$ time python3 p.py
15485867
python3 p.py 5.94s user 0.02s system 100% cpu 5.949 total
$ time pypy3 p.py
15485867
pypy3 p.py 2.52s user 0.06s system 100% cpu 2.585 total
Then just wrap the whole thing in a main function, because local
variables are faster than global ones:
from itertools import accumulate
from math import log
def main():
# ...
main()
Times:
$ time python3 p.py
15485867
python3 p.py 4.74s user 0.02s system 100% cpu 4.754 total
$ time pypy3 p.py
15485867
pypy3 p.py 2.40s user 0.06s system 100% cpu 2.456 total
But wait! We dropped the end point!
for i in range(1, len(marked)//3):
step = 2*i + 1
start = i+step-1
stop = 2*i*(i+1)
marked[start:stop:step] = [False] * len(slice_proxy[start:stop:step])
This is important on PyPy; on CPython it's about the same speed:
$ time python3 p.py
15485867
python3 p.py 4.88s user 0.01s system 100% cpu 4.890 total
$ time pypy3 p.py
15485867
pypy3 p.py 1.44s user 0.02s system 99% cpu 1.461 total
It's not really faster on CPython because the number of things you do tends to be more important
than what you do, simply from interpreter overhead.
We know it's way faster on CPython, but how does this compare to the original version on PyPy
(once stuck inside a main function)?
from math import log
def main():
n = 1000001
maximum = n * (log(n) + log(log(n)))
marked = [False] * int(maximum // 2)
for i in range(1, len(marked)+1):
for j in range(1, i+1):
m = i + j + 2*i*j
if m > len(marked):
break
marked[m-1] = True
count = 1
for i in range(1, len(marked)+1):
if not marked[i-1]:
count += 1
if count == n:
print(2*i + 1)
break
main()
Times:
$ time pypy p.py
15485867
pypy p.py 0.83s user 0.02s system 99% cpu 0.851 total
$ time pypy3 p.py
15485867
pypy3 p.py 0.78s user 0.01s system 99% cpu 0.790 total
Interestingly, the original is significantly faster on PyPy. This implies
that unless you want a code base optimized for both CPython and PyPy, some
of these optimizations (like slice_proxy) just aren't needed.
But, eh, what if this isn't fast enough?
Eh?
Well, let's quickly draw the things we're setting. T are the parts we set. . is the part we would set
if not for the :end: optimization.
index: 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
i=1 T . . . . . . . . .
i=2 T T . . . .
i=3 T T T .
i=4 T T T
i=5 T T
i=6 T T
i=7 T
i=8 T
i=9 T
i=10 T
i=11 T
We're doing i=1 then i=2 then i=3, etc. From i² ≈ len(marked) onwards, though, there aren't many
things set each run. We end up moving back-and-forth in the array a lot, spending more time "seeking" in
memory than actually setting items. Consider, though, this decomposition:
index: 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
i=1 1 2 3 4 5 6 7 8 9 10
i=2 1 2 3 4 5 6
i=3 1 2 3 4
i=4 1 2 3
i=5 1 2
i=6 1 2
i=7 1
i=8 1
i=9 1
i=10 1
i=11 1
For low numbers, 1, 2, 3, ... this hits a lot of those later values. We can then split the work into
an upper part:
i=1 T . . . . . . . . .
i=2 T T . . . .
i=3 T T T .
i=4 T T T
i=5 T T
and a lower part:
i=6 1 2
i=7 1
i=8 1
i=9 1
i=10 1
i=11 1
This sets the same values, but does so in a different order.
Here's the code:
from itertools import accumulate
from math import log
def main():
n = 1000001
maximum = n * (log(n) + log(log(n)))
marked = [True] * int(maximum // 2)
slice_proxy = range(len(marked))
for i in range(1, int(len(marked)**0.5)):
step = 2*i + 1
start = 3*i
stop = 2*i*(i+1)
marked[start:stop:step] = [False] * len(slice_proxy[start:stop:step])
i += 1
start_step = 2*i + 1
start_start = 3*i
step = 3
for start in range(start_start, len(marked), start_step):
marked[start::step] = [False] * len(slice_proxy[start::step])
step += 2
for i, count in enumerate(accumulate(marked), 1):
if count+1 == n:
print(2*i + 1)
break
main()
How fast is it?
$ time python3 p.py
15485867
python3 p.py 1.53s user 0.01s system 99% cpu 1.538 total
$ time pypy3 p.py
15485867
pypy3 p.py 1.26s user 0.06s system 100% cpu 1.319 total
CPython improves a lot with this, although PyPy doesn't really improve. This makes sense: you do a lot
fewer operations but each operation is more complicated. Since CPython has a really fast implementation
under the hood but a really slow interpreter, fewer operations is important. PyPy, however, is bound
by the size of the operations themselves.
Of course, for PyPy you'd want to modify the "original" code:
from math import log
def main():
n = 1000001
maximum = n * (log(n) + log(log(n)))
marked = [True] * int(maximum // 2)
for i in range(1, int(len(marked)**0.5)):
step = 2*i + 1
start = 3*i
stop = 2*i*(i+1)
if stop > len(marked):
stop = len(marked)
for m in range(start, stop, step):
marked[m] = False
i += 1
start_step = 2*i + 1
start_start = 3*i
step = 3
for start in range(start_start, len(marked), start_step):
for m in range(start, len(marked), step):
marked[m] = False
step += 2
count = 1
for i in range(1, len(marked)+1):
if marked[i-1]:
count += 1
if count == n:
print(2*i + 1)
break
main()
Which gives
$ time pypy p.py
15485867
pypy p.py 0.65s user 0.02s system 99% cpu 0.672 total
$ time pypy3 p.py
15485867
pypy3 p.py 0.61s user 0.01s system 99% cpu 0.625 total
The difference is obviously much less, but it exists. This is probably
due to improvements in memory locality. It actually turns out that just
terminating the range with len(marked)//3 gets you almost to this
point, so the optimization, IMHO, isn't worth it.
Sadly, CPython still trails the first version on PyPy. Janne Karila gave a good suggestion in the comments for speeding up CPython. The key part is to use itertools (or filter) to apply a filter and islice to do the counting. The large advantage comes from the iteration and counting being done inside C routines.
I went with compress from Python 3.1 and Janne Karila's use of islice to get the wanted element:
primes = compress(range(3, int(maximum), 2), marked)
print(next(islice(primes, n-2, None)))
Timings are really good:
$ time python3 p.py
15485867
python3 p.py 0.58s user 0.01s system 99% cpu 0.598 total
This is (negligibly) faster than PyPy!
But, eh, but... but...
What if this isn't fast enough?
Eh?
Well, that's what Numpy is for. Numpy is Python's fast array library.
from math import log
import numpy
def main():
n = 1000001
maximum = n * (log(n) + log(log(n)))
marked = numpy.ones(int(maximum // 2), dtype=bool)
for i in range(1, int(len(marked)**0.5)):
step = 2*i + 1
start = 3*i
stop = 2*i*(i+1)
marked[start:stop:step] = False
i += 1
start_step = 2*i + 1
start_start = 3*i
step = 3
for start in range(start_start, len(marked), start_step):
marked[start::step] = False
step += 2
accumulated = numpy.add.accumulate(marked)
[where] = numpy.where(accumulated+1 == n)
print(2 * where[0] + 3)
main()
This also works on Python 2. Unfortunately PyPy doesn't (yet)
have an easy Numpy port. Even so, the times:
$ time python3 p.py
15485867
python3 p.py 0.23s user 0.02s system 99% cpu 0.252 total
$ time python2 p.py
15485867
python2 p.py 0.17s user 0.02s system 98% cpu 0.196 total
C++ with this N takes:
$ g++ -O3 c.cpp -o c; time ./c
15485867
./c 0.06s user 0.00s system 96% cpu 0.065 total
Neat, but we're still not done. One slow part is
accumulated = numpy.add.accumulate(marked)
[where] = numpy.where(accumulated+1 == n)
print(2 * where[0] + 3)
It turns out that Numpy's count_nonzero is really fast, so let's try capitalizing on that with what amounts to a binary search. There are similarities to Quickselect:
# I can't think of a good name...
def count_nonzero_sorted(array, count):
# We always assume the count is achievable
if array.shape[0] == 1:
return 0
# Split the array in two
midpoint = len(array) // 2
first_half = array[:midpoint]
second_half = array[midpoint:]
first_counts = numpy.count_nonzero(first_half)
# If the counts in the first half is greater
# than the wanted index, check that half
if first_counts >= count:
return count_nonzero_sorted(first_half, count)
# else check the other half and offset the result
return midpoint + count_nonzero_sorted(second_half, count-first_counts)
This gives
# compress sieve
where = count_nonzero_sorted(marked, n-1)
print(2 * where + 3)
and timings of just
$ time python2 p.py
15485867
python2 p.py 0.10s user 0.01s system 97% cpu 0.116 total
$ time python3 p.py
15485867
python3 p.py 0.16s user 0.01s system 99% cpu 0.167 total
On Python 2 that's half the speed of C++! Vectorization is a wonderful thing.
It would be interesting how many of these optimizations can apply to R.
A preliminary guess puts the time required for the original R code to run as about 7 hours.
What should you take away from this?
PyPy is really, really fast. The original code (after the
bug is fixed) takes only 0.7-0.9 seconds, or a tenth the speed of C++.
PyPy thus ends up, by my guess, about 10000 times as fast
as the implementation of R that I have available (for this task).
Optimizing PyPy is possible but it's hard; the optimized PyPy was
only a little faster than the original PyPy code at 0.6-0.7s.
This is because the JIT gives you most optimizations for free.
CPython has sharp corners (range vs xrange) but Python 3 takes
some of these problems away by default. Once you handle this, CPython
is an order of magnitude faster than R on this code.
CPython can be forced to go fast if you know what you are doing.
When forced it became a touch faster than PyPy.
CPython is ridiculously fast when Numpy is used; getting to half the speed
of C++.
I thought to check how a traditional sieve compares against this:
import sys
from itertools import compress, islice
from math import log
def main():
n = 1000001
maximum = int(n * (log(n) + log(log(n))))
maxidx = maximum//2
sieve = [True] * maxidx # 2, 3, 5, 7... might be prime
j = 0
for i in range(1, int(maxidx**0.5)+1):
j += 4*i
if sieve[i]:
step = 2*i + 1
sieve[j::step] = [False] * -(-(maxidx-j) // step)
# compress sieve
print(next(islice(compress(range(1, maximum, 2), sieve), n-1, None)))
main()
and the PyPy version:
import sys
from itertools import compress, islice
from math import log
def main():
n = 1000001
maximum = int(n * (log(n) + log(log(n))))
maxidx = maximum//2
sieve = [True] * maxidx # 2, 3, 5, 7... might be prime
j = 0
for i in range(1, int(maxidx**0.5)+1):
j += 4*i
if sieve[i]:
for k in range(j, maxidx, 2*i + 1):
sieve[k] = False
# compress sieve
count = 1
for i in range(1, maxidx):
count += sieve[i]
if count == n:
print(2*i + 1)
break
main()
and the Numpy version:
import numpy
from math import log
def count_nonzero_sorted(array, count):
# We always assume the count is achievable
if array.shape[0] == 1:
return 0
midpoint = len(array) // 2
first_half = array[:midpoint]
second_half = array[midpoint:]
first_counts = numpy.count_nonzero(first_half)
if first_counts >= count:
return count_nonzero_sorted(first_half, count)
return midpoint + count_nonzero_sorted(second_half, count-first_counts)
def main():
n = 1000001
maximum = int(n * (log(n) + log(log(n))))
maxidx = maximum//2
sieve = numpy.ones(maxidx, dtype=bool) # 2, 3, 5, 7... might be prime
j = 0
for i in range(1, int(maxidx**0.5)+1):
j += 4*i
if sieve[i]:
sieve[j::2*i + 1] = False
# compress sieve
where = count_nonzero_sorted(sieve, n)
print(2 * where + 1)
main()
Timings:
$ time python3 p.py
15485867
python3 p.py 0.40s user 0.01s system 99% cpu 0.415 total
$ time pypy3 p.py
15485867
pypy3 p.py 0.36s user 0.01s system 99% cpu 0.368 total
$ # Numpy versions
$ time python2 p.py
15485867
python2 p.py 0.07s user 0.01s system 97% cpu 0.075 total
$ time python3 p.py
15485867
python3 p.py 0.12s user 0.00s system 99% cpu 0.124 total
Sorry, Sieve of Sundaram, you're not that great after all. :P
And, heck, that Python 2 version is way too fast.
Note: I made a mistake previously in comparing vector<bool> to vector<int>; I should have compared it to vector<uint8_t> instead. It turns out that the difference is actually relatively minor (30-40%) and although significant it isn't nearly the factor-of-4 difference from before.
