Universal Tree Resolver FTW

gistfile1.txt

#!/usr/bin/ruby1.9.1

results = [
           {id: 4, data: "level2", pid: 3},
           {id: 1, data: "level1", pid: 0},
           {id: 0, data: "level0", pid: nil},
           {id: 2, data: "level2", pid: 1},
           {id: 2, data: "levelX", pid: 10},
           {id: 3, data: "level1", pid: 0}
          ]

class TreeBuilder
  class Node
    attr_accessor :id, :pid, :data, :children

    def initialize(args)
      @id, @pid, @data = args[:id], args[:pid], args[:data]
      @children = []
    end

    def print(depth = 0)
      puts "#{"  " * depth}#{@data}"
      @children.each do |c|
        c.print(depth + 1)
      end
    end
  end

  def initialize
    @tree = nil
    @references = {}
    @orphins = []
  end
  
  def item(it)
    n = Node.new it

    if @tree.nil?
      @tree = n
      @references[it[:id]] = @tree
      return
    end

    if it[:pid].nil?
      oldtree = @tree
      @tree = n
      if oldtree.pid == @tree.id
        @tree.children << oldtree
      else
        @orphins << oldtree
      end
      @references[it[:id]] = @tree
      return
    end

    if @references[it[:pid]]
      @references[it[:pid]].children << n
      @references[it[:id]] = @references[it[:pid]].children.last
      return
    end

    @orphins << n
  end

  # Returns the best tree we can generate.
  def tree
    resolve_orphins!
    @tree
  end

  private

  # Iteratively resolve orphins.
  def resolve_orphins!
    loop do
      old_orphin_count = @orphins.count

      @orphins.map! do |o|
        if @references[o.pid]
          @references[o.pid].children << o
          @references[o.id] = @references[o.pid].children.last
          nil
        else
          o
        end
      end.compact!
      
      break if @orphins.count == old_orphin_count
    end
  end
end

tb = TreeBuilder.new

results.each do |r|
  tb.item r
end

tb.tree.print