First, put everything in a main
function. This makes the variables local, which makes lookups a little faster. It's also neater.
It's worth formatting the comments with line wrapping and putting them in front of the commented line to keep lines short.
collatz_array[1]
is never used except to iterate over, so just iterate over range(2, p)
instead and ditch collatz_array[1]
. In fact, at point of use it's simpler to just call enumerate
:
max_value, max_count = 0, 0
for i, val in enumerate(collatz_array[2], 2):
if max_count < val:
max_count = val
max_value = i
Then you should split collatz_array
into two arrays, which I'm tentatively calling progression_mask
and progression_length
. They can be initialized as
progression_mask = np.arange(2, p, dtype='int64')
collatz_array = np.zeros(p-2, dtype='int64')
count
is just len(sequence)
, so generate it at the end with len
. When looping over sequence
you can do:
count = len(sequence)
for q in sequence:
...
count -= 1
if q in progression_mask
is \$O(n)\$ but you can do the check in \$O(1)\$ with
if q < p and progression_mask[q-2]:
The same goes for if w in progression_mask
.
This gets you much, much faster... but we can go faster still.
I would remove the -2
from the indexes by changing the arrays to:
progression_mask = np.arange(p, dtype='int64')
progression_length = np.zeros(p, dtype='int64')
Using pure lists is faster than Numpy arrays, so switch to them:
progression_mask = list(range(p))
progression_length = [0] * p
The final loop can be simplified with a max
and index
:
max_count = max(progression_length)
max_value = progression_length.index(max_count)
The main loop would be better with an early exit on the if i != 0
check, and it should be spelt if i
(or if not i
when inverted).
for i in progression_mask:
if not i:
continue
Generating the collatz numbers should be done with a function:
def collatz_progression(n):
while n > 1:
yield n
if n % 2 == 0:
n = n//2
else:
n = int((3*n)+1)
Note that this doesn't yield 1
, as I don't think it's needed. It's called as
sequence = list(collatz_progression(i))
This isn't needed once we've moved to plain lists:
q = int(q)
If progression_mask
is 0
, we don't need to continue so can break from the loop:
if q < p:
if not progression_mask[q]:
break
...
count -= 1
I would ditch progression_mask
and just check if progression_length
is set.
Your loop
w = q
# Loop is run for about the number of even multiples of q
# the next even multiple
for e in range(1, p//q):
w *= 2
if w < p and progression_mask[w]:
would be better as
w = q
count_ = count
while w < p:
progression_length[w] = count_
count_ += 1
w *= 2
in part so it can replace the other assignment. I would move it into a new function and rename the variables:
def assign_ascending_collatz(start, count, progression_length):
while start < len(progression_length):
progression_length[start] = count
start *= 2
count += 1
n = int((3*n)+1)
can just be n = 3*n + 1
.
if not n % 2
is more idiomatic than if n % 2 == 0
, but even better is to flip the conditional and do
if n % 2:
n = 3*n + 1
else:
n = n//2
or maybe even
n = (3*n + 1) if n % 2 else (n // 2)
This gives:
def collatz_progression(n):
while n > 1:
yield n
n = (3*n + 1) if n % 2 else (n // 2)
def assign_ascending_collatz(start, count, progression_length):
while start < len(progression_length):
progression_length[start] = count
start *= 2
count += 1
def main():
up_to = 1000000
progression_length = [-1] * up_to
for i in range(up_to):
if progression_length[i] != -1:
continue
sequence = list(collatz_progression(i))
for i, val in enumerate(sequence):
if val < up_to and progression_length[val] != -1:
break
# Cache the count
assign_ascending_collatz(val, len(sequence)-i, progression_length)
max_count = max(progression_length)
max_value = progression_length.index(max_count)
print("Number with longest collatz sequence: ", max_value)
main()
With PyPy3 this takes me ~1.5s to run.
Most of the time seems to be in collatz_progression
, so we can speed that up by letting it short-circuit against progression_length
. This gives:
def collatz_progression(n, progression_length):
while True:
yield n
if n < len(progression_length) and progression_length[n] != -1:
return
n = (3*n + 1) if n % 2 else (n // 2)
def assign_ascending_collatz(start, count, progression_length):
while start < len(progression_length):
progression_length[start] = count
start *= 2
count += 1
def main():
up_to = 1000000
progression_length = [-1] * up_to
# Skip 0 and set 1
progression_length[0] = 0
progression_length[1] = 0
for i in range(up_to):
if progression_length[i] != -1:
continue
sequence = list(collatz_progression(i, progression_length))
count = len(sequence) + progression_length[sequence.pop()]
for val in sequence:
# Cache the count
assign_ascending_collatz(val, count, progression_length)
count -= 1
max_count = max(progression_length[1:])
max_value = progression_length.index(max_count)
print("Number with longest collatz sequence: ", max_value)
main()
Note that I had to fix progression_length[1]
to avoid the -1
adding to the other values.
With PyPy3 this takes ~0.3 seconds to run. Hopefully that's fast enough.
Note that now your clever optimization of going forward as well with assign_ascending_collatz
doesn't actually help. Getting rid of it even gives a tiny speed improvement:
def collatz_progression(n, progression_length):
while True:
yield n
if n < len(progression_length) and progression_length[n] != -1:
return
n = (3*n + 1) if n % 2 else (n // 2)
def main():
up_to = 1000000
progression_length = [-1] * up_to
# Skip 0 and set 1
progression_length[0] = 0
progression_length[1] = 0
for i in range(up_to):
if progression_length[i] != -1:
continue
sequence = list(collatz_progression(i, progression_length))
count = len(sequence) + progression_length[sequence.pop()]
for val in sequence:
if val < len(progression_length):
progression_length[val] = count
count -= 1
max_count = max(progression_length[1:])
max_value = progression_length.index(max_count)
print("Number with longest collatz sequence: ", max_value)
main()
Some good questions from the comments:
Before turning the lists into arrays, for some p, the execution time was 305s and after the change, it became 448s. Does it not count against lists?
Numpy arrays are faster than lists for some things and slower for others. In particular, arrays tend to be slower when indexing single items (eg. my_array[12]
). They tend to be faster on operations that affect the whole array, such as q in progression_mask
, which searches the whole array.
Before removing if q in progression_mask
and if w in progression_mask
, most of the time was spent in those two operations. This meant arrays were significantly faster.
After removing them, all operations on the array were just indexes. This means that lists became faster.
So lists are only faster after making certain changes to the code.
Why did you populate the progression_length with -1? It didn't seem to affect the result, either with 0 or -1.
I (arbitrarily) set progression_length[1] = 0
(instead of 1
). This meant that 0
was a valid value, so it didn't make sense in my opinion to use it as a sentinel.
I would have chosen None
instead of -1
if PyPy didn't have a special optimization for lists of only integers, meaning it's faster if I use an integer. If you don't want to use PyPy or don't mind PyPy being a little slower, I suggest using None
.
And you lost me completely after you 'short-circuited collatz_progression against progression_length'.
Let's say you have these lengths:
0 1 2 3 4 5 6 7
[ 0 0 1 6 1 4 ? ? ]
where the ?
is represented with -1
.
We want to find the collatz number for 6
. The original method would find this list:
[6, 3, 10, 5, 16, 8, 4, 2]
find its length and then loop over it, setting the values in progression_length
s.
However, once you have gotten
[6, 3, ...]
you know that every number afterwards is already set in the array since progression_length[3] != -1
. This means you don't need to set them in the array.
You still need to know the number of skipped items, so you look it up from
# The last value before we stopped
progression_length[3]
which gives 6
. You add that to the number of values in front (1
) and get 7
. This means you only have to generate 2 values from collatz_progression
instead of 7
in order to get the appropriate length.
And why didn't even multiples elimination not work?
It did work but the above optimization does the same thing in a better way.