Build a Tree of Structs

amedeo

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

2 comments

#functional-programming #structs #trees

1 2539 2

2017-11-17 20:41:17 UTC

Showing Posts 1 to 2

Show Best Posts
Show All (oldest first)
Show All (newest first)

michalmuskala

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}

Post #1