Collatz conjecture with plots

Question

I have an assignment where I write a Collatz Conjecture program for a series of starting values from 1 to N and make two plots: number of iterations vs starting value and computed numbers vs starting value. The assignment prompt is:

Now modify your program so that it computes 3 vectors. The first vector s(k) is the starting number (N), the vector f(k) stores the computed number as a function of starting number N at iteration k and g(N) is the number of iterations needed to get to the number 1. For example, if N=4, s, f and g would look like: 𝑠 = [1,2,2,3,3,3,3,3,3,3,3,4,4,4] 𝑓 = [1,2,1,3,10,5,16,8,4,2,1,4,2,1] 𝑔 = [0,1,7,2]

Make a plot of these values (in one window if possible), and make its x and y axis range from 1 to N. Use the number range from 1 to about 200 for both axes.

Here is my code.

import numpy as np
import matplotlib.pyplot as plt

# INPUTS
# Set N to a positive integer. The Hailstone problem program
# will be evaluated for starting numbers from 1 to N.
N = 200
marker_size = 7

s_no_repeats = range(1, N+1) # starting number vector (with NO repeating starting values)

def flatten(t):
    ''' Create flatlist out of list of lists '''
    return [item for sublist in t for item in sublist]

def find_repeat(numbers):
    '''Check repeating value in list and returns the value'''
    seen = set()
    for num in numbers:
        if num in seen:
            return num
        seen.add(num)

def Collatz(n):
    k = 0 # current iteration number
    computed_nums = [n] # sequence of computed numbers from n to 1
    iterations = [0]
    while n != 1:
        if n % 2 == 0:
            n = (n / 2)
        else:
            n = ((n * 3) + 1)
        k += 1
        computed_nums.append(n)
        iterations.append(k)
    
    map(int, computed_nums)
    return (computed_nums, k)

s = [] # starting number vector (with repeating starting values)
f = [] # computed number vector as function of starting number N at iteration k
g = [] # number of iterations required to get to 1

for starting_num in s_no_repeats:
    temp_f, temp_g = Collatz(starting_num)
    #print(temp_f)
    f.append(temp_f)
    g.append(temp_g)
    s.append([starting_num] * (temp_g + 1))

s = flatten(s)
f = flatten(f)

print("s vector (starting nums): " + str(s))
print("f vector (computed nums): " + str(f))
print("g vector (iterations): " + str(g))

fig, (ax1, ax2) = plt.subplots(1, 2)

# left subplot (for computed values)
ax1.scatter(s, f, color="blue", s=marker_size, clip_on=False, zorder = 10)
ax1.set_title('range of computed values during iteration')
ax1.set_xlabel('starting value')
ax1.set_xlim([1, N])
ax1.set_ylim([1, N])

# right subplot (for iterations)
ax2.scatter(s_no_repeats, g, color="red", s=marker_size, clip_on=False, zorder = 10)
ax2.set_title('number of iterations')
ax2.set_xlabel('starting value')
ax2.set_xlim([1, N])
ax2.set_ylim([1, N])

fig.tight_layout() # automatically adjusts enough space between subplots

I was wondering if this can be improved?

Answer 1

You don't use Numpy, so delete your import.

N is a parameter and not a constant - so refactor your code into functions that pass it around, removing your code from the global namespace.

Delete your flatten function and use list.extend instead of append to build your lists up.

find_repeat is not used to delete it.

Collatz should be lower-case by PEP8.

Consider using divmod to perform a combined division by 2 and modulus by 2.

Delete iterations from your collatz function since you don't use it.

Give more meaningful names to your s, f, g lists and your axis objects.

clip_on and zorder are just making a mess; delete them.

Add a __main__ guard.

There is no value in changing colour between your two subplots since they have separate axes. You would want your separate colours if you attempt to plot both datasets on the same axes.

Add PEP484 type hints.

Consider reducing the alpha-channel of your points to illustrate point density more easily.

Your map call's return was unused, so delete it.

Suggested

import matplotlib.pyplot as plt


def collatz(n: int) -> tuple[
    list[int],  # computed_nums
    int,        # current iteration number
]:
    k = 0                # current iteration number
    computed_nums = [n]  # sequence of computed numbers from n to 1
    while n != 1:
        half_n, odd = divmod(n, 2)
        if odd:
            n = n*3 + 1
        else:
            n = half_n
        k += 1
        computed_nums.append(n)

    return computed_nums, k


def collatz_all(n: int) -> tuple[
    list[int],  # starting
    list[int],  # computed
    list[int],  # iterations
    range,      # starting_no_repeats
]:
    starting_no_repeats = range(1, n + 1)  # starting number vector (with NO repeating starting values)

    starting = []  # starting number vector (with repeating starting values)
    computed = []  # computed number vector as function of starting number N at iteration k
    iterations = []  # number of iterations required to get to 1

    for starting_num in starting_no_repeats:
        temp_f, temp_g = collatz(starting_num)
        computed.extend(temp_f)
        iterations.append(temp_g)
        starting.extend([starting_num] * (temp_g + 1))

    return starting, computed, iterations, starting_no_repeats


def plot(starting, computed, iterations, starting_no_repeats, n: int) -> plt.Figure:
    fig, (ax_computed, ax_iterations) = plt.subplots(nrows=1, ncols=2)
    style = {
        's': 7,
        'alpha': 0.2,
    }

    # left subplot (for computed values)
    ax_computed.scatter(starting, computed, **style)
    ax_computed.set_title('range of computed values during iteration')
    ax_computed.set_xlabel('starting value')
    ax_computed.set_ylim((0, n))

    # right subplot (for iterations)
    ax_iterations.scatter(starting_no_repeats, iterations, **style)
    ax_iterations.set_title('number of iterations')
    ax_iterations.set_xlabel('starting value')
    ax_iterations.set_ylim((0, n))

    fig.tight_layout()  # automatically adjusts enough space between subplots
    return fig


def main() -> None:
    # INPUTS
    # Set N to a positive integer. The Hailstone problem program
    # will be evaluated for starting numbers from 1 to N.
    n = 200

    starting, computed, iterations, starting_no_repeats = collatz_all(n)
    print(f"s vector (starting nums): {starting}")
    print(f"f vector (computed nums): {computed}")
    print(f"g vector (iterations): {iterations}")

    plot(starting, computed, iterations, starting_no_repeats, n)
    plt.show()


if __name__ == '__main__':
    main()

This plot is in a large window so the transparency effect isn't very pronounced:

It becomes more important for a condensed window: