Skip to content

Instantly share code, notes, and snippets.

@outofmbufs
Created July 20, 2021 00:27
Show Gist options
  • Save outofmbufs/c652f15112b566a5ff33958332a01e14 to your computer and use it in GitHub Desktop.
Save outofmbufs/c652f15112b566a5ff33958332a01e14 to your computer and use it in GitHub Desktop.
Python Collatz Sequence as in OEIS A006577
#!/usr/bin/env python3
def collatz_sequence(n):
"""Generate the collatz sequence for n (see OEIS A006577)."""
c = int(n)
if c < 1 or c != n:
raise ValueError(f"Argument n ({n}) must be a positive integer")
while c != 1:
yield c
if (c % 2) == 0:
c = c // 2
else:
c = (3 * c) + 1
yield 1
def collatz_steps(n):
"""Returns # of collatz steps as defined in OEIS A006577."""
return len(list(collatz_sequence(n))) - 1
if __name__ == "__main__":
import unittest
class TestMethods(unittest.TestCase):
def test1(self):
# test data literally copy-pasted from https://oeis.org/A006577
testdata = (0, 1, 7, 2, 5, 8, 16, 3, 19, 6, 14, 9, 9, 17, 17,
4, 12, 20, 20, 7, 7, 15, 15, 10, 23, 10, 111, 18,
18, 18, 106, 5, 26, 13, 13, 21, 21, 21, 34, 8, 109,
8, 29, 16, 16, 16, 104)
for n, k in enumerate(testdata):
with self.subTest(n=n, k=k):
self.assertEqual(collatz_steps(n+1), k)
def test2(self):
badvals = (0, -1, 2.5)
for bad in badvals:
with self.subTest(bad=bad):
with self.assertRaises(ValueError):
next(collatz_sequence(bad))
def test3(self):
# arguably perhaps this should be an error
self.assertEqual(collatz_steps(4.0), 2)
unittest.main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment