Search code examples
pythondata-structuresbinary-tree

How to display a Binary Search Tree


I am trying to display a binary search tree in Python using the _displayRec method below. However, when I test it with a simple example, the display becomes unbalanced on the right side:

def display(self):
        print("Displaying tree:")
        print("----------------")
        lines, *_ = self._displayRec(self._root)
        for line in lines:
            print(line)
        print("----------------")

    def _displayRec(self, node):
        if node is None:
            line = "Empty"
            width = len(line)
            height = 1
            middle = width // 2
            return [line], width, height, middle
        else:
            key = str(node.getKey())
            value = str(node.getValue())
            label = f"{key} ({value})"
            left_lines, left_width, left_height, left_middle = self._displayRec(node.getLeft())
            right_lines, right_width, right_height, right_middle = self._displayRec(node.getRight())
            node_label = label
            node_label_width = len(node_label)
            first_line = (left_middle + 1) * " " + (left_width // 2) * "_" + node_label + (right_width // 2) * "_"
            second_line = left_middle * " " + "/" + (left_middle + node_label_width + right_middle) * " " + "\\" + (right_width - right_middle - 1) * " "
            if left_height < right_height:
                left_lines += [left_width * " "] * (right_height - left_height)
            elif right_height < left_height:
                right_lines += [right_width * " "] * (left_height - right_height)
            zipped_lines = zip(left_lines, right_lines)
            lines = [first_line, second_line] + [a + " " * (node_label_width + 1) + b for a, b in zipped_lines]
            width = left_width + node_label_width + right_width + 1
            height = max(left_height, right_height) + 2
            middle = width // 2
            return lines, width, height, middle

Here's the code I'm using to test the display method:

tree.insert(4, "four")
tree.insert(5, "five")
tree.insert(1, "one")
tree.insert(2, "two")
tree.insert(3, "three")
tree.display()

The resulting display looks like this:

Inserted key 4 with value four
Inserted key 5 with value five
Inserted key 1 with value one
Inserted key 2 with value two
Inserted key 3 with value three
Displaying tree:
----------------
                        _______________________4 (four)_________
                       /                                        \
   __1 (one)________________            __5 (five)__
  /                         \                           /            \
Empty           __2 (two)__________         Empty         Empty
               /                   \
             Empty           __3 (three)__
                            /             \
                          Empty          Empty
----------------

As you can see, the right side of the display is unbalanced and difficult to read. How can I modify the _displayRec method to fix this issue and make the tree display balanced?


Solution

  • One of the problems is that your lines are not equally sized. All of the following expressions should be the same value, but they are different:

    • len(first_line)
    • len(second_line)
    • left_width + node_label_width + 1 + right_width

    I took a slightly different approach where the returned lines include an extra line with the upwards edge "/", and then it is up to the caller to change that to a backslash or to throw it away (in case it is the initial call).

    The use of the .center() and .rjust() string methods makes some expressions easier.

    This version will not have the problem you encountered:

        def display(self):
            print("Displaying tree:")
            print("----------------")
            lines = self._displayRec(self._root)[1:]
            for line in lines:
                print(line)
            print("----------------")
    
        def _displayRec(self, node):
            if node is None:
                return [
                    "   /   ",
                    " Empty "
                ]
            else:
                left_lines = self._displayRec(node.getLeft())
                right_lines = self._displayRec(node.getRight())
                right_lines[0] = right_lines[0].replace("/", "\\")
    
                label = f"{node.getKey()} ({node.getValue()})"
                label_width = len(label)
                label_offset = left_lines[0].index("/") + 1
                available_label_width = len(left_lines[0]) + right_lines[0].index("\\") - label_offset
                gap = max(0, label_width - available_label_width)
                available_label_width += gap
                width = len(left_lines[0]) + gap + len(right_lines[0])
    
                diff = len(left_lines) - len(right_lines) 
                if diff < 0:
                    left_lines += [len(left_lines[0]) * " "] * (-diff)
                elif diff > 0:
                    right_lines += [len(right_lines[0]) * " "] * diff
                return [
                    (" " * label_offset + "/".center(available_label_width)).ljust(width),
                    (" " * label_offset + label.center(available_label_width, "_")).ljust(width)
                ] + [a + " " * gap + b for a, b in zip(left_lines, right_lines)]