How can I write Collatz conjecture by SymPy?

Question

I was going to write a SymPy function which takes a natural number and returns the step number of Collatz conjecture of the arg. The Python version is here:

def collatz(n: int):
    step = 0
    while n != 1:
        n = n // 2 if n % 2 == 0 else 3 * n + 1
        step += 1
    return step

print(collatz(27))

It prints 111. How about SymPy version? I feel sympy.Lambda() should have a recursive feature, similar to recursive call of procedual programming. Is there a good way?

• Just calling collatz() with a sympy.Symbol() instance (obviously) went into iloop.
• sympy.series.sequences.RecursiveSeq() does only backward reference with constant decrements.

Answer 1

Let f(x) be the symbolic collatz result. Do a substitution followed by a replacement to replace any non-symbolic results:

>>> f=Function('f')
>>> eq = f(x)
>>> eq.subs(x, 270)
f(270)
>>> _.replace(lambda x: x.func == f and x.args[0].is_Integer, lambda x: collatz(x.args[0]))
42

If you don't want to do it like this and want automatic evaluation, then you will have to write a SymPy class deriving from Function that has an eval method that detects when the input is an Integer. You can look at any function to see how this is implemented, e.g. see the source code for cos.