コラッツの問題の数列ジェネレーター

collatz_generator.py

def collatz_gen(n):
    yield n
    while n > 1:
        if n % 2 == 0:
            n = n // 2
        else:
            n = n * 3 + 1
        yield n

#print(list(collatz_gen(6)) == [6, 3, 10, 5, 16, 8, 4, 2, 1])
for i in collatz_gen(6):
    print(i)

# 1から9999までの数に対して、1にいたるまでのステップ数のグラフ

import matplotlib.pyplot as plt
import numpy as np

a = np.arange(1, 10000, 1)
b = [len(tuple(collatz_gen(i))) for i in a]

plt.scatter(a, b, s=10)
plt.show()