Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Collatz tree generator. Output as text or png (png requires graphviz). Produces regular binary collatz tree or trees without even numbers.
# Collatz tree generator. Output as text or png (png requires graphviz).
# Produces regular binary collatz tree or trees without even numbers.
#
# Note: The tree starts at 8 to avoid the 4-2-1 loop.
# Note': See the end of the file for two example text trees.
#
# Setup - on linux or wsl:
# ```shell
# sudo apt install python3-pip ipython3 graphviz
# sudo -H pip3 install anytree
# ```
# ipython3 isn't mandatory.
#
# Launch:
# ```
# $ ipython3 -i collatz.py
# >>> pngRellabeledTree()
# >>> # The above commands generates a file named
# >>> # "collatz-sophisticated-re-200.png" in the currend working directory.
# >>>
# >>> pngRellabeledTree(400, "Sophisticated")
# >>> # > "collatz-sophisticated-re-400.png"
# >>>
# >>> pngTree(150, "Binary")
# >>> # > "collatz-binary-150.png"
# >>>
# >>> printTree(100, "Binary")
# >>> # Example output at the end of the file (`print(example[0])`)
# >>>
# >>> printTree(200, "Sophisticated")
# >>> # Example output at the end of the file (`print(example[1])`)
# ```
from abc import abstractmethod
import anytree as a
from anytree import Node, RenderTree, AbstractStyle
from anytree.exporter import DotExporter
from sys import argv
mode = "Binary Sophisticated".split()[1]
# * Binary is the regular Collatz tree. Nodes have either 1 or 2 children.
# * Sophisticated is a Collatz tree without even numbers (except 8). Nodes have
# either 0 or infinite children.
try:
z = int(argv[1])
except Exception:
z = 200
def main():
pass
else:
def main():
printTree()
# Frontend
def printTree(z=None, mode=None):
"""
Produce text tree with labeled elements. See `pngRellabledTree` for the
explanation about labels.
"""
__printTree(z, mode)
def pngTree(z=None, mode=None):
"""
Produce a png image of the tree, using graphviz.
"""
__pngTree(z, mode)
def pngRellabeledTree(z=None, mode=None):
"""
Produce a png image of the tree, including labels '_/\'.
The labels indicate for each element, whether it is congruent
to 0 (_), 1 (/) or 2 (\) modulo 3.
"""
__pngRellabeledTree(z, mode)
# Backend
def __info(v, patterns="_/\\"):
return " " + patterns[v % 3]
def __label(node, info=__info):
return "{}{}".format(node.name, info(node.name))
semiStyle = AbstractStyle("│ ", "├─", "└─")
def __render(tree):
r = RenderTree(tree, style=semiStyle)
return "\n".join("{}{}".format(pre, __label(node)) for pre, _, node in r)
def __getClass(mode):
return globals()[f"CollatzTree{mode}"]
def __printTree(_z=None, _mode=None):
if _z is None:
_z = z
if _mode is None:
_mode = mode
tree = Node(8)
__getClass(_mode)(_z).addChildrenOf(tree)
print(__render(tree))
def __pngTree(_z=None, _mode=None):
tree = Node(8)
__getClass(_mode)(_z).addChildrenOf(tree)
DotExporter(tree).to_picture(f"collatz-{_mode.lower()}-{_z}.png")
def __relabelTree(tree):
reInfo = lambda v: __info(v, patterns=["_", "/", "\\\\"])
for _, _, node in tuple(RenderTree(tree)):
node.name = __label(node, info=reInfo)
def __pngRellabeledTree(_z=None, _mode=None):
if _z is None:
_z = z
if _mode is None:
_mode = mode
tree = Node(8)
__getClass(_mode)(_z).addChildrenOf(tree)
__relabelTree(tree)
DotExporter(tree).to_picture(f"collatz-{_mode.lower()}-re-{_z}.png")
class AbstractCollatzTree:
def __init__(self, z):
self.z = z
self.y = 3 * z + 1
@abstractmethod
def childrenOf(self, node):
pass
def addChildrenOf(self, node):
for child in self.childrenOf(node):
if child >= self.z:
break
try:
self.addChildrenOf(Node(child, parent=node))
except RecursionError:
import pdb
pdb.post_mortem()
class CollatzTreeBinary(AbstractCollatzTree):
"""
Binary Collatz Tree.
Contains both even and odd numbers.
"""
def childrenOf(self, node):
assert node not in [1, 2, 4]
if node.name % 6 == 4:
yield (node.name - 1) // 3
yield 2 * node.name
# Note that the two yields must be in increasing order
class CollatzTreeSophisticated(AbstractCollatzTree):
"""
Sophisticated Collatz Tree.
Contains only odd numbers. Each node theorically has an infinit number of
decendents.
"""
def childrenOf(self, node):
epsilon = node.name % 3
if epsilon == 0: return
child = (2 ** (3 - epsilon) * node.name - 1) // 3
while True:
yield child
child = 4 * child + 1
if __name__ == "__main__":
main()
example = {}
example[0] = example["""printTree(100, "Binary")"""] = r"""
8 \
└─16 /
├─5 \
│ └─10 /
│ ├─3 _
│ │ └─6 _
│ │ └─12 _
│ │ └─24 _
│ │ └─48 _
│ │ └─96 _
│ └─20 \
│ └─40 /
│ ├─13 /
│ │ └─26 \
│ │ └─52 /
│ │ └─17 \
│ │ └─34 /
│ │ ├─11 \
│ │ │ └─22 /
│ │ │ ├─7 /
│ │ │ │ └─14 \
│ │ │ │ └─28 /
│ │ │ │ ├─9 _
│ │ │ │ │ └─18 _
│ │ │ │ │ └─36 _
│ │ │ │ │ └─72 _
│ │ │ │ └─56 \
│ │ │ └─44 \
│ │ │ └─88 /
│ │ │ └─29 \
│ │ │ └─58 /
│ │ │ └─19 /
│ │ │ └─38 \
│ │ │ └─76 /
│ │ │ └─25 /
│ │ │ └─50 \
│ │ └─68 \
│ └─80 \
└─32 \
└─64 /
└─21 _
└─42 _
└─84 _
"""
example[1] = example["""printTree(200, "Sophisticated")"""] = r"""
8 \
├─5 \
│ ├─3 _
│ ├─13 /
│ │ ├─17 \
│ │ │ ├─11 \
│ │ │ │ ├─7 /
│ │ │ │ │ ├─9 _
│ │ │ │ │ ├─37 /
│ │ │ │ │ │ ├─49 /
│ │ │ │ │ │ │ └─65 \
│ │ │ │ │ │ │ ├─43 /
│ │ │ │ │ │ │ │ └─57 _
│ │ │ │ │ │ │ └─173 \
│ │ │ │ │ │ │ └─115 /
│ │ │ │ │ │ │ └─153 _
│ │ │ │ │ │ └─197 \
│ │ │ │ │ │ └─131 \
│ │ │ │ │ │ └─87 _
│ │ │ │ │ └─149 \
│ │ │ │ │ └─99 _
│ │ │ │ ├─29 \
│ │ │ │ │ ├─19 /
│ │ │ │ │ │ ├─25 /
│ │ │ │ │ │ │ ├─33 _
│ │ │ │ │ │ │ └─133 /
│ │ │ │ │ │ │ └─177 _
│ │ │ │ │ │ └─101 \
│ │ │ │ │ │ └─67 /
│ │ │ │ │ │ └─89 \
│ │ │ │ │ │ └─59 \
│ │ │ │ │ │ ├─39 _
│ │ │ │ │ │ └─157 /
│ │ │ │ │ └─77 \
│ │ │ │ │ └─51 _
│ │ │ │ └─117 _
│ │ │ ├─45 _
│ │ │ └─181 /
│ │ └─69 _
│ └─53 \
│ ├─35 \
│ │ ├─23 \
│ │ │ ├─15 _
│ │ │ └─61 /
│ │ │ └─81 _
│ │ └─93 _
│ └─141 _
├─21 _
└─85 /
└─113 \
└─75 _
"""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.