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 =  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?