Plotting a binary tree in matplotlib

Question

Suppose I have a binary tree of depth d, represented by a array of length 2^d - 1. Where left_child(i)=2*i + 1, right_child(i)=2*i + 2 I want to plot the tree to get something like the following full binary tree of depth 4

(The blue lines can be ignored) The figure above was plotted with a matrix of shape (depth, 2^depth - 1) where the first column entries all had the root's value (0 in this case), the second column had the upper half as the root's right child and the bottom half as its left child and so on until the last column which had all of the leaf values

I there a simpler way of getting this plot from the linear representation of the tree? Preferably with keeping different branches in different colors

JohanC · Accepted Answer · 2019-11-29 16:06:19Z

Here is some code to draw a binary tree. The function recursively draws the left and right subtrees. You can play with the number of levels and the internal distances. If you exchange x and y, you get a vertical layout. The colors are cycled from the currently active colormap.

import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib import colors as mcolors

width_dist = 10
depth_dist = 10
levels = 5

def bintree_level(levels, x, y, width):
    segments = []
    xl = x + depth_dist
    yl = y - width / 2
    xr = x + depth_dist
    yr = y + width / 2
    segments.append([[x, y], [xl, yl]])
    segments.append([[x, y], [xr, yr]])
    if levels > 1:
        segments += bintree_level(levels - 1, xl, yl, width / 2)
        segments += bintree_level(levels - 1, xr, yr, width  / 2)
    return segments

segs = bintree_level(levels, 0, 0, width_dist)

colors = [mcolors.to_rgba(c)
          for c in plt.rcParams['axes.prop_cycle'].by_key()['color']]
line_segments = LineCollection(segs, linewidths=1, colors=colors, linestyle='solid')

fig, ax = plt.subplots()
ax.set_xlim(-1, levels * depth_dist + 1)
ax.set_ylim(-1.5*width_dist, 1.5*width_dist)
ax.add_collection(line_segments)
plt.show()

Collectives™ on Stack Overflow

Plotting a binary tree in matplotlib

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related