Why flatten a tree when you can just traverse it?
Feb 15, 2013 · 1 minute read · 0 CommentsprogrammingRubyrecursion
Michael Feathers wrote a blog post about avoiding explicitly traversing a tree recursively by taking advantage of its representation as nested arrays and using array operations to flatten.
I was unhappy about this solution because
- it creates a bunch of intermediate arrays
- it concatenates arrays repeatedly, which is expensive for a very large tree
- the recursive solution seems more simple, clear, and efficient
So I wrote up the recursive solution. Some features of the solution:
- it uses Ruby’s idiomatic
yield
construct - for convenience, it creates an
Enumerable
so that you can do whatever you want while iterating through the tree
Here’s the code:
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
require 'minitest/autorun' | |
# Traverse in preorder | |
def tree_traverse(tree, &block) | |
if tree.is_a? Array and tree.length > 0 | |
yield tree | |
tree.each {|t| tree_traverse(t, &block) } | |
end | |
end | |
class TreeEnumerable | |
include Enumerable | |
def initialize(tree) | |
@tree = tree | |
end | |
def each | |
tree_traverse(@tree) {|t| yield t } | |
end | |
end | |
def sexp_select(tree, symbols) | |
TreeEnumerable.new(tree).find_all {|e| symbols.include? e[0] } | |
end | |
class TestTreeSelect < MiniTest::Unit::TestCase | |
def test_find_instance_variables | |
test_tree = [:program, [[:class, [:const_ref, [:@const, "A", [1, 6]]], nil, | |
[:bodystmt, [[:def, [:@ident, "a", [1, 13]], [:params, nil, nil, nil, nil, nil], | |
[:bodystmt, | |
[[:assign, | |
[:var_field, [:@ivar, "@a", [1, 16]]], | |
[:var_ref, [:@ivar, "@b", [1, 21]]]]], nil, nil, nil]]], nil, nil, nil]]]] | |
result = [[:@ivar, "@a", [1, 16]], [:@ivar, "@b", [1, 21]]] | |
assert_equal result, sexp_select(test_tree, [:@ivar]) | |
end | |
end |