diff --git a/lib/mb_binary_tree/traverse.ex b/lib/mb_binary_tree/traverse.ex index 97e771a..8dedf97 100644 --- a/lib/mb_binary_tree/traverse.ex +++ b/lib/mb_binary_tree/traverse.ex @@ -2,4 +2,53 @@ defmodule MbBinaryTree.Traverse do @moduledoc """ Utility functions for traversing binary trees. """ + + @doc """ + Perform a preorder traversal of a binary tree. + + ## Examples + + iex> MbBinaryTree.Traverse.preorder([1, [2, nil, nil], [3, nil, nil]]) |> Enum.to_list() + [1, 2, 3] + iex> MbBinaryTree.Traverse.preorder([12, [4, MbBinaryTree.new(2), MbBinaryTree.new(5)], [26, MbBinaryTree.new(14), [28, MbBinaryTree.new(27), MbBinaryTree.new(30)]]]) |> Enum.to_list() + [12, 4, 2, 5, 26, 14, 28, 27, 30] + """ + def preorder(tree) do + Stream.unfold([tree], &preorder_acc/1) + end + + defp preorder_acc([tree | stack]) do + case tree do + [value | [left, right]] -> {value, [left, right | stack]} + nil -> preorder_acc(stack) + end + end + + defp preorder_acc([]), do: nil + + @doc """ + Perform a postoder traversal of a binary tree. + + ## Examples + + iex> MbBinaryTree.Traverse.postorder([1, [2, nil, nil], [3, nil, nil]]) |> Enum.to_list() + [2, 3, 1] + iex> MbBinaryTree.Traverse.postorder([12, [4, MbBinaryTree.new(2), MbBinaryTree.new(5)], [26, MbBinaryTree.new(14), [28, MbBinaryTree.new(27), MbBinaryTree.new(30)]]]) |> Enum.to_list() + [2, 5, 4, 14, 27, 30, 28, 26, 12] + """ + def postorder(tree) do + Stream.unfold([tree], &postorder_acc/1) + end + + defp postorder_acc([tree | stack]) do + case tree do + [value | [nil, nil]] -> {value, stack} + [value | [left, nil]] -> postorder_acc([left, [value, nil, nil] | stack]) + [value | [nil, right]] -> postorder_acc([right, [value, nil, nil] | stack]) + [value | [left, right]] -> postorder_acc([left, right, [value, nil, nil] | stack]) + nil -> postorder_acc(stack) + end + end + + defp postorder_acc([]), do: nil end