Search code examples
functional-programmingtreeelixirheap

How to balance a Max Heap Tree in Elixir?


I've implemented a Max Heap Tree, but when creating new nodes, the tree becomes unbalanced. For example, if most of the values inserted are less than the root value, it becomes left-heavy tree. This is because of the if-else comparison, but is there any other way to balance the tree? Thanks in advance.

defmodule MaxHeap do

  defstruct value: nil, right: nil, left: nil

  def new(value), do: %MaxHeap{value: value}

  def insert(newValue=%MaxHeap{value: rootValue}, nextValue) do
    if nextValue <= rootValue do
      %MaxHeap{newValue | left: insert(newValue.left, nextValue)}
    else
      temp=rootValue #save old rootValue in temp, before editing it
      rootValue=nextValue
      nextValue=temp
      %MaxHeap{newValue | right: insert(newValue.right, nextValue), value: rootValue}
    end
  end

  def insert(nil, value) do
    %MaxHeap{value: value}
  end
end

Output tree in Elixir iex:


Solution

  • If what you really want is a heap, you can use a slightly different data structure called a leftist tree - it's a tree that's always unbalanced in a particular way, that allows you to have O(log(N)) heap operations. You can see a description of that in this blog post about leftist trees.

    One strategy to balance your tree that is simple to understand is what's called an AVL tree. An AVL tree is a binary tree where the difference in height between the two children of any node cannot be greater than one. This tree is almost balanced in that its height can be at most 2log(N) at any time. The way to achieve it here is to:

    1. Store the tree height in the nodes:

      defstruct value: nil, right: nil, left: nil, height: 1
      
    2. Update the tree height after insertion and then balance it:

      def insert(newValue = %MaxHeap{value: rootValue}, nextValue) do
        if nextValue <= rootValue do
          %MaxHeap{newValue | left: insert(newValue.left, nextValue)}
        else
          %MaxHeap{newValue | right: insert(newValue.right, nextValue)}
        end
        |> set_height()
        |> balance()
      end
      
    3. Have a height function that handles empty trees:

      def height(nil), do: 0
      def height(tree), do: tree.height
      
    4. Have a set_height function that sets the height of the parent based on the heights of the children:

      defp set_height(tree) do
        %{tree | height: max(height(tree.left), height(tree.right)) + 1}
      end
      
    5. In the balance function apply a rotation left or right if the subtrees vary by more than 1 in height:

      defp balance(tree) do
        cond do
          height(tree.left) > height(tree.right) + 1 -> rotate_right(tree)
          height(tree.right) > height(tree.left) + 1 -> rotate_left(tree)
          true -> tree
        end
      end
      
    6. Have a rotate_left function (and an analogous rotate_right function):

      defp rotate_right(tree) do
        %{tree.left | right: set_height(%{tree | left: tree.left.right})}
        |> set_height()
      end
      
      defp rotate_left(tree) do
        %{tree.right | left: set_height(%{tree | right: tree.right.left})}
        |> set_height()
      end
      

    The important thing to note about these rotations is that they preserve the binary tree invariant.

    You can read more about this approach in this geeksforgeeks post about AVL trees.