Skip to content

Instantly share code, notes, and snippets.

@mapio

mapio/Printing trees.ipynb

Last active January 9, 2019 11:34
Show Gist options
  • Save mapio/37c97e3fac25b59bf6b729b6c9446230 to your computer and use it in GitHub Desktop.
Save mapio/37c97e3fac25b59bf6b729b6c9446230 to your computer and use it in GitHub Desktop.
Download ZIP
Some recursive code to print trees
Display the source blob
Display the rendered blob
Raw
Printing trees.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Printing arbitrary trees\n",
"\n",
"This gist presents two ways to print (text representation of) *trees*. \n",
"\n",
"## The basics, representing and traversing trees\n",
"\n",
"First of all, to make tbings simple, we assume to represent a (rooted and integer labelled) tree using Python lists. The first element of the list is the root label, and the other elements of the list are the subtrees. So, for instance, a simple single node tree labelled $9$ is"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"simple = [9]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"whereas, just as an example, a binary tree labelled with integers where (if present) the left subtree of $n$ is labeled with $2n$ and the right one with $2n + 1$ and the root is labbeled with $1$ is"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"binary = [\n",
" 1, \n",
" [2, \n",
" [4, \n",
" [8], \n",
" [9]\n",
" ], \n",
" [5]\n",
" ], \n",
" [3, \n",
" [6], \n",
" [7]\n",
" ]\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Given a visitor function `visit`, the usual recursive *in-order* and *post-order* traversal of trees represented in this way are"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def preorder(tree, visit):\n",
" visit(tree[0])\n",
" for s in tree[1:]: preorder(s, visit)\n",
" \n",
"def postorder(tree, visit):\n",
" for s in tree[1:]: postorder(s, visit) \n",
" visit(tree[0]) "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"It is easy to convince oneself of the correctenss of both simply using `print` as a visitor"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1\n",
"2\n",
"4\n",
"8\n",
"9\n",
"5\n",
"3\n",
"6\n",
"7\n"
]
}
],
"source": [
"preorder(binary, print)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"8\n",
"9\n",
"4\n",
"5\n",
"2\n",
"6\n",
"7\n",
"3\n",
"1\n"
]
}
],
"source": [
"postorder(binary, print)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"A bit more complex is the *level-order*, as usual it requires a *queue* that can be implemented (as [suggested in the Python tutorial](https://docs.python.org/3/tutorial/datastructures.html?highlight=queue#using-lists-as-queues)) with a `dequeue`."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"from collections import deque\n",
"\n",
"def levelorder(tree, visit):\n",
" Q = deque()\n",
" Q.appendleft(tree)\n",
" while Q:\n",
" tree = Q.pop()\n",
" visit(tree[0])\n",
" for s in tree[1:]: Q.appendleft(s)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Again using `print` as a visit function can convince ourself of the correctenss of `levelorder`."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1\n",
"2\n",
"3\n",
"4\n",
"5\n",
"6\n",
"7\n",
"8\n",
"9\n"
]
}
],
"source": [
"levelorder(binary, print)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Let's pretty print now\n",
"\n",
"Suppose now that we want to print such trees in more \"readable\" ways that we have done using `print` and one of the above traversals.\n",
"\n",
"Our first goal could be to at least obtain a representation similar to the (handcrafted) one used in the cell above that defines `binary`. To such aim, we need to keep track of the depth of a node in the tree, the deepst the node, more spaces we should prepend when printing his label.\n",
"\n",
"This can be done simply adapting a preorder traversing to include the `depth` information. There is no need to store it along with the node, given that it is known at the time in which `print` is called."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"def prettyprint(tree):\n",
" def _prettyprint(tree, depth):\n",
" print(' ' * depth, tree[0])\n",
" for s in tree[1:]: _prettyprint(s, depth + 1)\n",
" _prettyprint(tree, 0)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" 1\n",
" 2\n",
" 4\n",
" 8\n",
" 9\n",
" 5\n",
" 3\n",
" 6\n",
" 7\n"
]
}
],
"source": [
"prettyprint(binary)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If we want a more conventional shape where every subtree appears below its root, we could the level-order traversal. \n",
"\n",
"We need at least to keep track of the end of every level (to start a new line). To this aim we put a *sentinel* (here we use `None`) in the queue to mark the end of one level; if we pop the sentinel we start a new line and, if the queue is not empty, we move the sentinel to the end of the queue."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"def levelprint(tree):\n",
" Q = deque()\n",
" Q.appendleft(tree)\n",
" Q.appendleft(None)\n",
" while Q:\n",
" tree = Q.pop()\n",
" if tree is None:\n",
" print()\n",
" if Q: Q.appendleft(None)\n",
" else:\n",
" print(tree[0], end = '')\n",
" for s in tree[1:]: Q.appendleft(s)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Again, printing the `binary` tree should give us some confidence in such implementation."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1\n",
"23\n",
"4567\n",
"89\n"
]
}
],
"source": [
"levelprint(binary)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We got right the \"vertical\" arrangement, but we completely miss the \"horizontal\" spacing, we should aim at something like\n",
"<pre>\n",
" 1\n",
" 2 3\n",
" 4 5 6 7\n",
" 8 9\n",
"</pre>\n",
"where `8 9` appear below `4`, `6 7` below `3` and the subtrees rooted `4` and `5` appear below `2`.\n",
"\n",
"To compute the horizontal position one can keep track of how many leafs have been ecountered during a post-order traversal, annotating such number as the *column* where the node should appear.\n",
"\n",
"To this aim, let's extend the node from a simple integer to a [dataclass](https://docs.python.org/3/library/dataclasses.html) cotaining the node label (as a string) and column (observe in passing that [namedtuples](https://docs.python.org/2/library/collections.html#collections.namedtuple) can't be used because their fields are *immutable*)."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"from dataclasses import make_dataclass, field\n",
"\n",
"Node = make_dataclass('Node', [\n",
" ('label', str),\n",
" ('col', int, field(default = 0)),\n",
"])\n",
"\n",
"str_binary = [\n",
" Node('one'), \n",
" [Node('two'), \n",
" [Node('four'), \n",
" [Node('eight')], \n",
" [Node('nine')]\n",
" ], \n",
" [Node('five')]\n",
" ], \n",
" [Node('three'), \n",
" [Node('six')], \n",
" [Node('seven')]\n",
" ]\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Given that we now use strings as labels, to compute the column we should conisder the label length, plus eventually some separator."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"SEP = ' '\n",
"\n",
"def setcol(tree):\n",
" col = 0\n",
" def _setcol(tree):\n",
" nonlocal col\n",
" tree[0].col = col\n",
" if len(tree) == 1: col += len(SEP + tree[0].label)\n",
" for s in tree[1:]: _setcol(s)\n",
" _setcol(tree)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's apply this to the binary tree with string labels"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Node(label='one', col=0),\n",
" [Node(label='two', col=0),\n",
" [Node(label='four', col=0),\n",
" [Node(label='eight', col=0)],\n",
" [Node(label='nine', col=6)]],\n",
" [Node(label='five', col=11)]],\n",
" [Node(label='three', col=16),\n",
" [Node(label='six', col=16)],\n",
" [Node(label='seven', col=20)]]]"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"setcol(str_binary)\n",
"str_binary"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We now should be done, we can improve `levelprint` to work with string labels and to keep track of the current column (adding spaces to what is about to print in order to match the node column)."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [],
"source": [
"def betterprint(tree):\n",
" Q = deque()\n",
" Q.appendleft(tree)\n",
" Q.appendleft(None)\n",
" col = 0\n",
" while Q:\n",
" tree = Q.pop()\n",
" if tree is None:\n",
" print()\n",
" col = 0\n",
" if Q: Q.appendleft(None)\n",
" else:\n",
" tp = (tree[0].col - col) * ' ' + SEP + tree[0].label\n",
" print(tp, end = '', sep = '')\n",
" col += len(tp)\n",
" for s in tree[1:]: Q.appendleft(s)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Running the code seems to lead to what we where looking for."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" one\n",
" two three\n",
" four five six seven\n",
" eight nine\n"
]
}
],
"source": [
"betterprint(str_binary)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.1"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment