class TrieStoreTest < Test::Unit::TestCase
  def setup
    @t = TrieStore.new
  end
 
  def test_should_insert
    @t.insert('hello', 1)
  end
 
  def test_should_search
    @t.insert('hello', :value1)
    assert_equal [ :value1 ], @t.search('hello')
  end
 
  def test_should_search_with_prefix
    @t.insert('hello', :value1)
    assert_equal [ :value1 ], @t.search('he')
  end
 
  def test_should_not_search_missing_values
    @t.insert('hello', :value)
    assert !@t.search('goodbye')
  end
 
  def test_should_return_unique_matches
    @t.insert('hello', :value1)
    @t.insert('help', :value2)
    assert_equal [ :value1 ], @t.search('hello')
  end
 
  def test_should_return_common_matches
    @t.insert('hello', :value1)
    @t.insert('help', :value2)
    assert_equal [ :value1, :value2 ], @t.search('hel')
  end
end

class TrieStore
  def initialize
    @children, @data = {}, []
  end
 
  def insert(term, value)
    return if term.empty?
 
    head = term[0, 1]
    rest = term[1..-1]
    @data |= [value]
    @children[head] ||= TrieStore.new
    @children[head].insert(rest, value)
  end
 
  def search(term)
    case term.size
      when 0 then []
      when 1 then @data
      else
        head = term[0, 1]
        rest = term[1..-1]
        if child = @children[head]
          child.search(rest)
        else
          []
        end
    end
  end
end

Leave a Comment

Enclose code in <code lang="ruby"></code> if you care.
Preview your comment using the button below.