Build a Tree of Structs

Questions / Help
, ,
1

Hi,
I’m working on a Elixir/Phoenix small project in my spare time and I found myself writing code to build a tree structure starting from a list of nodes.

Here is the input:

[
%{id: 1, children: [], parent_id: nil},
%{id: 2, children: [3], parent_id: nil},
%{id: 3, children: [], parent_id: 2}
]

and the expected output:

%TreeNode{
   id: "ROOT",
   children: [
      %TreeNode{
           id: 1,
           parent_id: "ROOT"
           children: []
      },
      %TreeNode{
           id: 2,
           parent_id: "ROOT"
           children: [
               %TreeNode{
                   id: 3,
                   parent_id: 2
                   children: []
                }
           ]
      }
  ],
  parent_id: nil
}

I’ve come to the following implementation but it is quite long and I’ve the feeling that it could be written in a different way.
Do you have some idea or suggestions on how to write it?

defmodule TestTree do

  defmodule TreeNode do
    defstruct id: nil, children: [], parent_id: nil
  end

  @_ROOT_NODE "ROOT"

  def build_tree(list) do
    # use the build_hierarchy to get two Maps as support data for the next step
    {id_to_nodes, id_to_children} = build_hierarchy(list, %{@_ROOT_NODE=>%TreeNode{id: @_ROOT_NODE}}, %{@_ROOT_NODE=>[]})
    # build the tree
    build_tree(Map.keys(id_to_nodes), id_to_nodes, id_to_children)
  end

  # Creates two Maps:
  # one Map keeps the information on Id -> %TreeNode
  # the other Map is an Id -> [ /child id 1/, /child id 2/, ...]
  defp build_hierarchy([], a, b), do: {a, b}
  defp build_hierarchy([m | nodes], id_to_nodes, id_to_children) do
    parent_key = m.parent_id || @_ROOT_NODE

    new_node = %TreeNode{ id: m.id, parent_id: parent_key}

    id_to_nodes = Map.put_new(id_to_nodes, m.id, new_node)

    # update the parent children
    id_to_children = Map.put_new_lazy(id_to_children, parent_key, fn -> [] end)
    children = id_to_children[parent_key]
    id_to_children = %{id_to_children | parent_key => [new_node.id | children]}

    build_hierarchy(nodes, id_to_nodes, id_to_children)
  end


  # the function will build the tree (tail) recursively examining each node_id
  # if the current node has no children then it is added to the parent node
  # if the node has children then its processing is postponed, we need first to take care of the ones without children
  defp build_tree([@_ROOT_NODE], id_to_nodes, _), do: id_to_nodes[@_ROOT_NODE]
  defp build_tree([node_id | node_ids], id_to_nodes, id_to_children) do
    node = id_to_nodes[node_id]
    parent_id = node.parent_id
    children = Map.get(id_to_children, node_id, [])
    if length(children) > 0 do
      # node not ready yet, it has some children
      build_tree(node_ids ++ [node_id], id_to_nodes, id_to_children)
    else
      # the node has no children let's add it to its parent children and remove it from the id_to_children
      parent_node = id_to_nodes[parent_id]
      id_to_nodes = %{id_to_nodes | parent_id => %{parent_node | children: [node | parent_node.children]}}
      id_to_children = %{id_to_children | parent_id => List.delete(id_to_children[parent_id], node_id)}
      build_tree(node_ids, id_to_nodes, id_to_children)
    end
  end

end
1 Like
2

That depends largely on the data. If the data is sorted (as in the sample you provided), it can be transformed using the following:

def build_tree(elements) do
  {tree, []} = build_tree(elements, %TreeNode{id: nil, children: []})
  tree
end

defp build_tree([%{parent_id: parent, id: id, children: _} | rest], 
                %{id: parent, children: children} = tree) do
  {node, rest} = build_tree(rest, %TreeNode{id: id, children: []})
  build_tree(rest, %{tree | children: [node | children]})
end
defp build_tree(rest, tree), do: {tree, rest}
7 Likes
3

:open_mouth: awesome! and beautiful.

Sry, I forgot about the ordering, I have it on the parent_id, nil firsts, but I’ll do a pass of sorting before calling the new build_tree :smiley:

Thanks!

1 Like