Improve collatz_sequence algorithm

maths/collatz_sequence.py

@@ -1,19 +1,33 @@
-def collatz_sequence(n):
+from typing import List
+
+
+def collatz_sequence(n: int) -> List[int]:
     """
-    Collatz conjecture: start with any positive integer n.Next term is obtained from the previous term as follows: 
-    if the previous term is even, the next term is one half of the previous term. 
-    If the previous term is odd, the next term is 3 times the previous term plus 1. 
-    The conjecture states the sequence will always reach 1 regaardless of starting value n.
+    Collatz conjecture: start with any positive integer n. The next term is
+    obtained as follows:
+        If n term is even, the next term is: n / 2 .
+        If n is odd, the next term is: 3 * n + 1.
+
+    The conjecture states the sequence will always reach 1 for any starting value n.
     Example:
+    >>> collatz_sequence(2.1)
+    Traceback (most recent call last):
+        ...
+    Exception: Sequence only defined for natural numbers
+    >>> collatz_sequence(0)
+    Traceback (most recent call last):
+        ...
+    Exception: Sequence only defined for natural numbers
     >>> collatz_sequence(43)
     [43, 130, 65, 196, 98, 49, 148, 74, 37, 112, 56, 28, 14, 7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]
     """
+
+    if not isinstance(n, int) or n < 1:
+        raise Exception("Sequence only defined for natural numbers")
+
     sequence = [n]
     while n != 1:
-        if n % 2 == 0:  # even number condition
-            n //= 2
-        else:
-            n = 3 * n + 1
+        n = 3 * n + 1 if n & 1 else n // 2
         sequence.append(n)
     return sequence
 
@@ -22,7 +36,7 @@ def main():
     n = 43
     sequence = collatz_sequence(n)
     print(sequence)
-    print("collatz sequence from %d took %d steps." % (n, len(sequence)))
+    print(f"collatz sequence from {n} took {len(sequence)} steps.")
 
 
 if __name__ == "__main__":