3n + 1 programming challenge in Python

Question

I am trying to find a efficient solution for the 3n + 1 problem on uvaonlinejudge. The code I have uses memoization using a dictionary. I am getting a 'Time limit Exceeded' error when I submit the code. Can anyone suggest an improvement(s) that will help with the execution time of this code.

I have already seen this post and implemented the proposed improvements. All other posts related to this are either for C++ or Java.

import sys

def recCycleLength(n,cycLenDict):
    if n==1:
        return 1
    if n not in cycLenDict:
        if n%2==0:
            cycLen = recCycleLength(n//2, cycLenDict)
            cycLenDict[n] = cycLen + 1
            return cycLen+1
        else:
            cycLen = recCycleLength(3*n+1, cycLenDict)
            cycLenDict[n] = cycLen + 1
            return cycLen+1
    else:
        return cycLenDict[n]


def maxCycle(a, b):
    i = a
    mydict = {} 
    maxLength = 1
    while i <= b:
        m = recCycleLength(i, mydict)
        if m > maxLength:
            maxLength = m
        i = i + 1
    return maxLength


for line in sys.stdin:
    curr_line=line.split()
    num1 = int(curr_line[0])
    num2 = int(curr_line[1])
    if num1>num2:
        num1, num2 = num2, num1
    m = maxCycle(num1, num2)
    print("{} {} {}".format(num1, num2, m))

Answer 1

You should separate your responsibilities. Currently your function recCycleLength (which should be named rec_cycle_length to conform with Python's official style guide PEP8) does two things, calculating the return value and caching it in a passed dictionary. It would be better to outsource the latter to a decorator, like this one (which is the fastest memoization decorator for a single-valued function I know of):

def memodict(f):
    """ Memoization decorator for a function taking a single argument """
    class memodict(dict):
        def __missing__(self, key):
            ret = self[key] = f(key)
            return ret
    return memodict().__getitem__


@memodict
def rec_cycle_length(n): 
    if n == 1: 
        return 1
    elif n % 2 == 0: 
        return cyc_len = rec_cycle_length(n//2) + 1
    return rec_cycle_length(3*n+1) + 1

You can even use this decorator on maxCycle when changing it to accept only one argument, the tuple (a, b). You could also be using max with a generator expression here:

@memodict
def max_cycle(bounds):
    a, b = bounds
    return max(rec_cycle_length(n) for n in range(a, b + 1))

(Use xrange in case you are using python 2.x.)

Your code on module level should be guarded by a if __name__ == '__main__': guard. You could also use map and sorted to get rid of quite a few lines of code:

if __name__ == '__main__':
    for line in sys.stdin:
        bounds = tuple(sorted(map(int, line.split())))
        print(" ".join(bounds), max_cycle(bounds))

This allows you to do e.g. from filename import rec_cycle_length in another script to reuse a function defined here.

Note that even though str.format is awesome, it is not needed here.

---

Since you seem to be using python 3.x, you might want to use functools.lru_cache instead of this decorator, like so:

import functools

@functools.lru_cache(2**10)
def rec_cycle_length(n):
    ...

At least on my machine it is a lot faster for some use cases (because the cache will not grow without bounds, making it better when there are few repetitions, i.e. many cache misses.) It is, however, only included in python >= 3.3.

Answer 2

Your performance problem is that you throw away the memoization cache after each test case.

Also, in Python, iteration should be more efficient than recursion.

Answer 3

after changes suggested by 200_success, here is the code that got accepted:

import sys

def recCycleLength(n,cycLenDict):
    if n==1:
        return 1
    if n not in cycLenDict:
        if n%2==0:
            cycLen = recCycleLength(n//2, cycLenDict)
            cycLenDict[n] = cycLen + 1
            return cycLen+1
        else:
            cycLen = recCycleLength(3*n+1, cycLenDict)
            cycLenDict[n] = cycLen + 1
            return cycLen+1
    else:
        return cycLenDict[n]


def maxCycle(a, b, mydict):
    i = a
    maxLength = 1
    while i <= b:
        m = recCycleLength(i, mydict)
        if m > maxLength:
            maxLength = m
        i = i + 1
    return maxLength

memDict = {1:1}
for line in sys.stdin:
    curr_line=line.split()
    num1 = int(curr_line[0])
    num2 = int(curr_line[1])
    if num1>num2:
        num1, num2 = num2, num1
    m = maxCycle(num1, num2, memDict)
    print("{} {} {}".format(num1, num2, m))

Answer 4

After changes suggested by 200_success. The code that accepts the 3n + 1 problem on uvaonlinejudge:

Made some changes to accept the problem on the platform uvaonlinejudge. The code answered was giving Wrong Answer. Based on OP comment, he was right and I rearranged the output:

import sys

def recCycleLength(n,cycLenDict):
    if n==1:
        return 1
    if n not in cycLenDict:
        if n%2==0:
            cycLen = recCycleLength(n//2, cycLenDict)
            cycLenDict[n] = cycLen + 1
            return cycLen+1
        else:
            cycLen = recCycleLength(3*n+1, cycLenDict)
            cycLenDict[n] = cycLen + 1
            return cycLen+1
    else:
        return cycLenDict[n]


def maxCycle(a, b, mydict):
    i = a
    maxLength = 1
    while i <= b:
        m = recCycleLength(i, mydict)
        if m > maxLength:
            maxLength = m
        i = i + 1
    return maxLength

memDict = {1: 1}
array = []
value = False
for line in sys.stdin:
    curr_line=line.split()
    num1 = int(curr_line[0])
    num2 = int(curr_line[1])
    if num1>num2:
        num1, num2 = num2, num1
        value = True
    m = maxCycle(num1, num2, memDict)
    
    if value:
        array.append("{} {} {}".format(num2, num1, m))
    else:
        array.append("{} {} {}".format(num1, num2, m))
    value = False
for i in array:
    print(i)