barinek.com: the official website

Simple Trie for Autocomplete

require 'test/unit'

class Trie
  def initialize()
    @root = TrieNode.new
    @entry_count = 0
  end

  def empty?
    @entry_count == 0;
  end

  def size
    @entry_count
  end

  def add(key, value)
    return if key.nil? || key.empty?

    current_node = @root
    key.each_char { |character|
      next_node = current_node.get(character)
      next_node = current_node.add(character) unless next_node
      current_node = next_node
    }
    current_node
    current_node.value = value
    @entry_count += 1
  end

  def contains(str)
    # todo...
  end

  def find(string)
    results = []
    (0..string.length-1).each { |start_index|
      start_index
      current_node = @root
      (start_index..string.length-1).each { |i|
        ch = string[i].chr
        current_node = current_node.get(ch)
        break if current_node.nil?
        results.push current_node.values if current_node.values?
      }
    }
    results
  end

  class TrieNode
    def initialize()
      @children = Hash.new
    end

    def value=(value)
      if (@values.nil?)
        @values = Array.new
      end
      @values.push value
    end

    def get(character)
      @children[character]
    end

    def add(character)
      node = TrieNode.new
      @children[character] = node
      node
    end

    def values
      @values
    end

    def values?
      !@values.nil?
    end
  end
end

class TriTest  Test::Unit::TestCase
  def test_empty?
    trie = Trie.new
    assert(trie.empty?)

    trie.add("key", "value")
    assert(!trie.empty?)
  end

  def test_size
    trie = Trie.new
    trie.add("key", "value")
    assert_equal(1, trie.size)
  end

  def test_ignores_nil_or_empty_string
    trie = Trie.new
    trie.add("", "value")
    trie.add(nil, "value")
    assert(trie.empty?)
  end

  def test_find
    trie = Trie.new
    trie.add("key", "value")
    assert_equal(1, trie.find("key").size)

    trie.add("keys", "value")
    assert_equal(2, trie.find("keys").size)
  end

  def test_allows_duplicates
    trie = Trie.new
    trie.add("key", "value")
    trie.add("key", "value")
    assert(!trie.empty?)
  end
end