Drawing sequentially represented binary trees in Python

I’ve been playing around with sequentially represented binary trees recently, i.e. a binary tree given as a list of integers where the node values are in contiguous sections of the list.

For example, a binary tree where the node values are also their indexes would be sequentially represented like this:

0 | 1, 2 | 3, 4, 5, 6 | 7, 8, 9, 10, 11, 12, 13, 14

Each section in the list represents one level of the binary tree, and has length $2^{l}$ where $l$ is the level. At the root, $l=0$ and $2^0=1$, so we know this section has a length of 1. The next section represents $l=1$, so has length $2^1=2$, and so on. At the fourth level, $l=3$, so there are $2^3=8$ nodes in that section.

While playing around with this data structure, I kept wanting to be able to quickly visualise what a given set of sequentially represented nodes would look like, and that turned into making a little python function to draw that in mono-spaced ascii characters.

The function signature is like this:

def draw_seq_bin_tree(nodes):
  pass

If we call it like this with sequential integers as the node values:

draw_seq_bin_tree(range(31))

Then a basic ascii drawing could look like this:

······························································
······························00······························
······························································
··············01······························02··············
······························································
······03··············04··············05··············06······
······························································
··07······08······09······10······11······12······13······14··
······························································
15··16··17··18··19··20··21··22··23··24··25··26··27··28··29··30
······························································

Notice that each “cell” in the drawing is 2 characters wide. This is to allow for integers up to 99 with consistent widths.

Basic sequential tree drawing

The basic problems to solve for that are:

  • Knowing the max level width so that we can space the preceding rows.
  • Knowing how much padding to put between nodes in the same level.
  • Placing parent nodes in the mid point between their child nodes.
  • Empty spacer rows between each level of nodes.

We can calculate the number of levels as $\lceil log_2{n+1} \rceil$, or in Python:

ceil(log(len(nodes) + 1, 2))

This allows for the last level to only be partially filled with nodes. E.g. if we’re given 10 nodes, the 4th level ends up with only 3 out of 8 possible nodes filled. That’s handled by the ceiling:

$\lceil log_2{n+1} \rceil = \lceil log_2{11} \rceil = \lceil 3.46 \rceil = 4$

So we’ll need 4 levels to be able to draw 10 nodes.

With the number of levels, we can then calculate the maximum level width with $2^{h-1}$ where $h$ is the height or number of levels. In Python:

2 ** (h - 1)

So with 4 levels, we have:

$2^{4-1} = 2^{3} = 8$

The 4th and final row will have 8 node slots, even if not all of them are filled with actual node values.

The next thing is to figure out how much spacing we need between each node on the same level.

We can figure that out visually by looking at the number of cells between each node. In this diagram the nodes are drawn as solid blocks so that the spacing widths are clearer:

······························································
······························██······························
······························································
··············██·············←15→·············██··············
······························································
······██·····←07→·····██·····←07→·····██·····←07→·····██······
······························································
··██·←03→·██·←03→·██·←03→·██·←03→·██·←03→·██·←03→·██·←03→·██··
······························································

Remember that each cell is 2 characters wide. The sequence 15, 7, 3 looks suspiciously related to descending powers of 2. So to get this inter-node spacing we can use $2^{h+1-l}-1$, where $h$ is the total number of levels and $l$ is the level we’re looking at, starting at 0 for the root node.

This means that at level index 2 in this tree we need spacing of

$2^{h+1-l}-1 = 2^{4+1-2}-1 = 2^3-1 = 8-1 = 7$

So we’ll put 7 spacer cells between each node in level index 2, as in the spacing diagram above.

This then means that the “box width” of nodes in a level is $j + (s \times (j - 1))$, where $j$ is the number of nodes and $s$ is the number of spacer cells in each gap. For example in level index 2, we have 4 nodes with 3 gaps of 7 spacer cells each between them, for a total box width in cells of:

$j + (s \times (j - 1)) = 4 + (7 \times (4 - 1)) = 25$

That means that the box width of the widest 4th level with 8 nodes will have a spacer width of

$2^{h+1-l}-1 = 2^{4+1-3}-1 = 3$

and a box width of

$n + (s \times (n - 1)) = 8 + (3 \times (8 - 1)) = 25$

The spacer gap width of the final level will always be 3 as we’re working up from that minimum level of spacing as we look back up the tree to the root node.

With that “max box width” and the box width of each level, we can centre each level horizontally by adding a margin of $(m - b) \div 2$ on either side, where $m$ is the max box width and $b$ is the box width of the level we’re looking at. It would be nice to include an extra margin of one cell around the whole drawing, so we end up with $((m - b) \div 2) + 1$ for the margin either side of each row.

This will automatically solve the placement of the parent nodes on the mid point of their child nodes due to the calculation of the spacer widths.

One tweak we need on that is to fill any empty slots and spacers on the last level, which might only be partially filled with nodes. For example with 10 nodes, the basic drawing would like this:

······························································
······························00······························
······························································
··············01······························02··············
······························································
······03··············04··············05··············06······
······························································
··07······08······09··········································
······························································

We know that the full number of slots in the tree

$2^h-1 = 2^4-1 = 15$

where $h$ is the height of the tree

We can then calculate the number of missing nodes in the last level as

$2^h-1-n = 2^4-1-10 = 15-10 = 5$

where $h$ is the height of the tree and $n$ is the actual number of nodes given, so we’ve got 5 empty slots and corresponding spacer gaps to fill with whitespace, which will be $5 + (5 \times 3) = 20$ cells as we have gaps on either side of these empty slots.

The final piece of this basic drawing is to add extra spacing rows between each level. We can do that by multiplying the number of levels by 2 and adding 1, and then drawing an empty row on even numbered rows.

Putting all of the above together for the basic drawing function, we have a Python function like this:

from math import ceil, log


def draw_seq_bin_tree(nodes):
  cell_width = 2
  empty_cell = "·" * cell_width
  n = len(nodes)
  nodes = [str(node).zfill(cell_width) for node in nodes]
  tree_height = ceil(log(n + 1, 2))
  max_lvl_nodes = 2 ** (tree_height - 1)
  last_lvl_gap = 3
  max_box_width = max_lvl_nodes + ((max_lvl_nodes - 1) * last_lvl_gap)
  row_width = max_box_width + 2
  for row_i in range((tree_height*2)+1):
    if row_i % 2 == 0:
      print(empty_cell * row_width)
    else:
      lvl_i = row_i // 2
      lvl_slots_n = 2 ** lvl_i
      lvl_start = lvl_slots_n - 1
      lvl_stop = (lvl_start * 2) + 1
      lvl_nodes = nodes[lvl_start:lvl_stop]
      lvl_gap_w = (2 ** (tree_height + 1 - lvl_i)) - 1
      lvl_box_w = (lvl_slots_n + (lvl_gap_w * (lvl_slots_n - 1)))
      lvl_margin_w = ((max_box_width - lvl_box_w) // 2) + 1
      lvl_margin = empty_cell*lvl_margin_w
      lvl_gap = empty_cell*lvl_gap_w
      lvl_fill = ""
      if lvl_i+1 == tree_height:
        max_slots = (2 ** tree_height) - 1
        missing_n = max_slots - n
        lvl_fill = empty_cell * (missing_n + (missing_n * lvl_gap_w))
      print(f"{lvl_margin}{lvl_gap.join(lvl_nodes)}{lvl_fill}{lvl_margin}")

Then we can draw ascii diagrams of trees like these:

··············
······00······
··············
··01······02··
··············
······························
··············00··············
······························
······01··············02······
······························
··03······04······05······06··
······························
······························································
······························00······························
······························································
··············01······························02··············
······························································
······03··············04··············05··············06······
······························································
··07······08······09······10······11··························
······························································

Adding level numbers

It would be helpful to be able to see the level index as a prefix to each level in the diagram. We can add that on the level rows, and add an extra filler on the spacer rows to account for it:

# ...
row_width = max_box_width + 3
# ...
lvl_legend = str(lvl_i).zfill(cell_width)
print(f"{lvl_legend}{lvl_margin}{lvl_gap.join(lvl_nodes)}{lvl_fill}{lvl_margin}")

Then we get a drawing like this:

································
00··············00··············
································
01······01··············02······
································
02··03······04······05······06··
································

Adding branch lines

Another helpful feature in the drawing would be lines indicating the branches from parent to child nodes. If we draw these as ·╱ and ╲· then we can place them on the position above each child node in the spacer row above.

The logic for placing branches like that is almost the same as the logic for placing the nodes themselves, so the easiest way to do it might be to just generate that row as normal as if it contained cells, and then swap the cells for a left or right branch alternately. The changes are:

# ...
ws = "·"
empty_cell = ws * cell_width
lft_br = empty_cell[:-1]+"╱"
rgt_br = "╲"+empty_cell[1:]
# ...
if row_i == 0 or row_i == rows_n-1:
  print(empty_cell * row_width)
else:
  # ... 
  lvl_nodes = nodes[lvl_start:lvl_stop]
  if row_i % 2 == 0:
    lvl_nodes = [lft_br if br % 2 == 0 else rgt_br for br in range(len(lvl_nodes))]

Adding node borders

The final thing that might enhance the drawing is to draw a box around each node using the box drawing characters.

We’d want each node to look something like this:

┌──┐
│01│
└──┘

That complicates the spacing logic a little bit, so the final (messy!) Python implementation looks like this:

from math import ceil, log


def draw_seq_bin_tree(nodes, cell_width=2):
    node_width = cell_width + 2
    ws = "·"
    empty_cell = ws * cell_width
    lft_br = f"{(ws * (node_width-2))}╱{ws}"
    rgt_br = f"{ws}╲{(ws * (node_width-2))}"
    cap = f"┌{'─' * cell_width}┐"
    base = f"└{'─' * cell_width}┘"
    space_width = 2

    n = len(nodes)
    nodes = [str(node).zfill(cell_width) for node in nodes]
    tree_height = ceil(log(n + 1, 2))
    max_lvl_nodes = 2 ** (tree_height - 1)
    max_lvl_gaps = max_lvl_nodes - 1
    last_lvl_gap_w = node_width

    max_box_width = (max_lvl_nodes * node_width) + (max_lvl_gaps * last_lvl_gap_w)
    row_width = max_box_width + (space_width * 3)
    rows_n = (tree_height * 4) + 1
    empty_row = ws * (row_width+(cell_width))

    for row_i in range(rows_n):
        if row_i == 0 or row_i == rows_n - 1:
            print(empty_row)
        else:
            lvl_i = row_i // 4
            lvl_slots_n = 2 ** lvl_i
            lvl_gaps_n = lvl_slots_n - 1
            lvl_start = lvl_slots_n - 1
            lvl_stop = (lvl_start * 2) + 1

            lvl_gap_w = (2 ** (tree_height + 2 - lvl_i)) - node_width
            lvl_box_w = (lvl_slots_n * node_width) + (lvl_gaps_n * lvl_gap_w)
            lvl_margin_w = ((row_width - lvl_box_w) // 2) + 1
            lvl_margin = ws * lvl_margin_w
            lvl_gap = ws * lvl_gap_w

            lvl_fill = ""
            if lvl_i + 1 == tree_height:
                max_slots = (2 ** tree_height) - 1
                missing_n = max_slots - n
                lvl_fill_w = (missing_n * node_width) + (missing_n * lvl_gap_w)
                lvl_fill = ws * lvl_fill_w

            lvl_nodes = [f"│{node}│" for node in nodes[lvl_start:lvl_stop]]
            lvl_legend = empty_cell

            if row_i % 4 == 0:
                # branch row
                lvl_nodes = [lft_br if br % 2 == 0 else rgt_br for br in range(len(lvl_nodes))]
            elif (row_i + 1) % 4 == 0:
                # base row
                lvl_nodes = [base for _ in lvl_nodes]
            elif (row_i + 3) % 4 == 0:
                # cap row
                lvl_nodes = [cap for _ in lvl_nodes]
            else:
                lvl_legend = str(lvl_i).zfill(cell_width)

            print(f"{lvl_legend}{lvl_margin[cell_width:]}{lvl_gap.join(lvl_nodes)}{lvl_fill}{lvl_margin}")

And it can produce diagrams like this:

····················
········┌──┐········
00······│00│········
········└──┘········
······╱······╲······
····┌──┐····┌──┐····
01··│01│····│02│····
····└──┘····└──┘····
····················
····································
················┌──┐················
00··············│00│················
················└──┘················
··········╱··············╲··········
········┌──┐············┌──┐········
01······│01│············│02│········
········└──┘············└──┘········
······╱······╲········╱······╲······
····┌──┐····┌──┐····┌──┐····┌──┐····
02··│03│····│04│····│05│····│06│····
····└──┘····└──┘····└──┘····└──┘····
····································
····································································
································┌──┐································
00······························│00│································
································└──┘································
··················╱······························╲··················
················┌──┐····························┌──┐················
01··············│01│····························│02│················
················└──┘····························└──┘················
··········╱··············╲················╱··············╲··········
········┌──┐············┌──┐············┌──┐············┌──┐········
02······│03│············│04│············│05│············│06│········
········└──┘············└──┘············└──┘············└──┘········
······╱······╲········╱······╲········╱······╲········╱······╲······
····┌──┐····┌──┐····┌──┐····┌──┐····┌──┐····┌──┐····┌──┐····┌──┐····
03··│07│····│08│····│09│····│10│····│11│····│12│····│13│····│14│····
····└──┘····└──┘····└──┘····└──┘····└──┘····└──┘····└──┘····└──┘····
····································································

View post: Drawing sequentially represented binary trees in Python