This is a read-only archive!

Ruby: fun with Trees

I had to come up with some kind of dependency tree at work recently. A bunch of files need to processed in a certain order. The current way this is done is by giving each file a rank number, 0-100 or so. Lower ranked files come first. All the files are read and then processed in order by rank. Files with equal rank can be done in any order.

This has some problems:

  1. Some files have rank 0. If I want to run a new file before those files, I need to change this rank from 0 to something higher. If there are already other files with rank 1, then I need to change all those. And so on and so forth.
  2. The script that does the processing is written in Python. Eww.

In reworking this problem, I realized that every file really only needs to come after either zero or one certain other files (and if this is not the case, I can force-fix things so this is the case). That means I have a tree where every node has one parent, and an arbitrary number of children. And I can process the files "in order" by doing a pre-order traversal of the tree. Computer science classes really do come in handy once in a while, I guess.

So I had the fun of writing a Tree class in Ruby. And I really was fun. It's a bit specialized due to how my processing script has to work, and I'm still debugging it, but this is most of it.

class Node
  attr_reader :value, :parent
  def initialize(value)
    @value = value
    @children = []
    @parent = nil
  end

  def add_child(child)
    raise "Invalid child type" unless child.is_a?(Node)
    @children.push(child)
    child.parent = self
  end

  def each_child
    @children.each do |child|
      yield child
    end
  end

  def ==(other)
    return @value == other.value
  end

  # Node can have only one parent; 
  #assigning a new one "moves" the node
  def parent=(node)
    @parent = node
  end

  def to_s
    @value.to_s
  end
end


class Tree
  attr_reader :root
  def initialize(root=nil)
    if root then
      if ! root.is_a?(Node) then
        @root = Node.new(root)
      else
        @root = value
      end
    end
  end

  def root=(value)
    if ! value.is_a?(Node) then
      @root = Node.new(value)
    else
      @root = value
    end
  end

  def insert(wanted, new_node)
    return false unless @root
    if ! new_node.is_a?(Node) then
      new_node = Node.new(new_node)
    end
    if ! wanted.is_a?(Node) then
      wanted = Node.new(wanted)
    end
    add(@root, wanted, new_node)
  end

  def add(node, wanted, new_node)
    #puts "Testing [#{node.value}] == [#{wanted}]"
    if node == wanted then
      node.add_child(new_node)
      return true
    else
      node.each_child do |child|
        return true if add(child, wanted, new_node)
      end
    end
    return false
  end

  def preorder(node=nil, &block)
    node ||= @root
    # stop processing if the block returns false
    if yield node then
      node.each_child do |child|
        preorder(child, &block)
      end
    end
  end
end

It was written in about an hour, so I'm still working on it. But it works, and it's a lot more elegant than the old way of doing things. For example I can do this:

tree.preorder do |node|
  puts "I'm looking at node #{node}!"
  node.each_child do |child|
    puts "#{child} is a child of #{node}"
  end
  do_something_interesting_with(node)
end
October 04, 2006 @ 7:20 AM PDT
Cateogory: Programming
Tags: Ruby