Using List comprehensions to solve Collatz conjecture

Viewed 173

Is there a way to use list comprehension's to workout Collatz conjecture without using a while statement or another method to append the n value to the ls without adding ls after each statement?

from random import choice
from time import sleep

n = choice([x for x in range(2, 99*99) if all(x%y != 0 for y in range(2, x))])
ls = []
ls.append(n)
while True:
    if n % 2 == 0:
        n = n // 2
        ls.append(n)
    elif n % 2 != 0:
        n = (3 * n) + 1
        ls.append(n)
    if n == 1:
        break
print(ls)
4 Answers

One way ("missing" the initial number, but I don't think that's important for the purpose):

print(f'{n}:')
print([n := 3*n+1 if n%2 else n//2
       for _ in iter(lambda: n, 1)])

Output for n = 92:

92:
[46, 23, 70, 35, 106, 53, 160, 80, 40, 20, 10, 5, 16, 8, 4, 2, 1]

And one (ab)using a list comp, based on kolypto's:

print([memo
       for memo in [[n]]
       for n in memo
       if n == 1 or memo.append(n//2 if n%2==0 else n*3+1)
      ][0])

Output for n = 92:

[92, 46, 23, 70, 35, 106, 53, 160, 80, 40, 20, 10, 5, 16, 8, 4, 2, 1]

I still think the while loop is the appropriate way, but it turns out it's not a lot faster. Times for solving all n from 1 to 5000:

 78 ms  while_loop_ewz93
 87 ms  list_comp_Kelly
126 ms  list_comp_kolypto
 82 ms  list_comp_kolypto_Kellied

My iter(lambda: n, 1) simulates while n != 1:. More generally, while condition: can be simulated with iter(lambda: bool(condition), False). (The explicit bool isn't necessary if the condition already is a bool, for example iter(lambda: mystring.startswith('x'), False).)

Benchmark code (Try it online!) with ewz93's and kolypto's modified to also not include the start number (for fairer comparison):

from timeit import repeat

def while_loop_ewz93(n):
    ls = []
    while n != 1:
        n = n // 2 if n % 2 == 0 else (3 * n) + 1
        ls.append(n)
    return ls

def list_comp_Kelly(n):
    return [n := 3*n+1 if n%2 else n//2
            for x in iter(lambda: n, 1)]

def list_comp_kolypto(n):
    return [
        *(lambda memo: [
            memo.append(n // 2 if n%2==0 else n*3+1) or memo[-1]
            for n in memo
            if memo[-1] != 1
        ])([n])
    ]

def list_comp_kolypto_Kellied(n):
    return [
        memo
        for memo in [[n]]
        for n in memo
        if n == 1 or memo.append(n//2 if n%2==0 else n*3+1)
    ][0]

funcs = [
    while_loop_ewz93,
    list_comp_Kelly,
    list_comp_kolypto,
    list_comp_kolypto_Kellied,
]

for _ in range(3):
    for func in funcs:
        t = min(repeat(lambda: list(map(func, range(1, 5001))), number=1))
        print('%3d ms ' % (t * 1e3), func.__name__)
    print()

Well while is what you use when you don't know yet how many steps it will take and since that is kind of baked into the logic of finding the conjecture for a value there is not really a way around it. I personally think there is nothing bad about using while loops.

You can still make the code a bit more compact and readable while keeping the while loop, e.g. like this:

from random import choice

n = choice([x for x in range(2, 99 * 99) if all(x % y != 0 for y in range(2, x))])

ls = [n]
while n != 1:
    n = n // 2 if n % 2 == 0 else (3 * n) + 1
    ls.append(n)
print(ls)

Edit:

A slightly modified version of @Kelly Bundys answer does the trick for me to make it even more compact (let's not mention the readability though):

from random import choice

n = choice([x for x in range(2, 99 * 99) if all(x % y != 0 for y in range(2, x))])

ls = [n] + [n := n // 2 if n % 2 == 0 else (3 * n) + 1 for _ in iter(lambda: n, 1)]
print(ls)
    seq = [
        *(lambda memo: [n] + [
            memo.append(n // 2 if n%2==0 else n*3+1) or memo[-1]
            for n in memo
            if memo[-1] != 1
        ])([n])
    ]

It's not a pure comprehension, but sort of.

The core idea was to introduce a named list, memo, which I can refer to as a variable and push values to

n=1024
ls=[n]
[ls.append(n//2 if n%2==0 else 3*n+1) for n in ls if n!=1]
print (ls)
Related