CoCalc Shared Fileswww / wiki / attachments / 10(2f)480b / collatz.py
Author: William A. Stein
1#!/usr/bin/env python
2"""
3Simple functions for experimenting with the Collatz sequence.
4
5>>> collatz(5)
6[5, 16, 8, 4, 2, 1]
7"""
8
9def collatz(n):
10    # base case: we're done if n is 1
11    if n == 1:
12        return [1]
13
14    # otherwise, we need to find the next element in the sequence.
15    #   - if n is even, it's n/2
16    #   - if n is odd, it's 3n+1
17    if n%2 == 0:
18        next_entry = n/2
19    else:
20        next_entry = 3*n+1
21
22    return [n] + collatz(next_entry)
23
24if __name__ == '__main__':
25    import sys
26    if len(sys.argv) != 2:
27        print "Incorrect number of arguments (expected 2, got %s)."%len(sys.argv)
28        exit(1)
29
30    ls = collatz(int(sys.argv[1]))
31    print "Sequence reached 1 after %s steps."%(len(ls)-1)
32    print " Entries:", ls
33
34def short_collatz(n):
35    import math
36    return (n == 1 and [1]) or \
37           ([n] + short_collatz(int(abs(math.sin(math.pi*n/2)))*(3*n+1) +
38                                int(abs(math.cos(math.pi*n/2)))*(n/2)))
39
40