Some functions for exploring the Collatz graph

collatz.py

from typing import Dict

import numpy as np

collatz_graph: Dict[int, int] = dict()
collatz_path_lengths: Dict[int, int] = dict()

def collatz(n: int):
    collatz_path_lengths[n] = collatz_iter(n)
    return collatz_path(n)

def collatz_path(n: int):
    path_length = collatz_path_lengths[n]
    path = np.zeros(path_length, dtype=np.int64)
    n_i = n
    path[0] = n
    for i in range(1, path_length):
        path[i] = collatz_graph[n_i]
        n_i = collatz_graph[n_i]
    return path

def collatz_iter(n: int):
    if n == 1:
        return 1

    n_i = collatz_graph.get(n, None)

    if n_i is None:
        if n % 2:
            n_i = (3*n) + 1
            collatz_graph[n] = n_i
        else:
            n_i = n / 2
            collatz_graph[n] = n_i

    return 1+collatz_iter(n_i)

print(collatz_graph)
print(collatz(97))
print(collatz_graph)