Searching words in dictionary

competition.rb

# The implementation of Trie is from https://github.com/khy/trie
class Trie
  include Enumerable
  
  attr_accessor :size
  
  def initialize
    @children = {}
    @terminal = false
    @size = 0
  end
  
  def add(key)
    visit_node(key) do |node, remainder|
      if remainder.empty?
        if !node.terminal?
          node.terminal!
          @size += 1
        end
      else
        node.append(remainder)
        @size += 1
      end

      self
    end
  end
  alias_method :<<, :add

  def delete(key)
    visit_node(key[0..-2]) do |node, remainder|
      if remainder.empty? and node.prune(key[-1])
        @size -= 1
        key
      end
    end
  end  
  
  def include?(key)
    visit_node(key) do |node, remainder|
      remainder.empty? and node.terminal?
    end
  end

  def prefixed_by(prefix)
    visit_node(prefix) do |node, remainder|
      remainder.empty? ? node.to_array(prefix) : []
    end
  end

  def each
    yield '' if terminal?

    @children.keys.sort.each do |value|
      @children[value].each{|key| yield Trie.join(value, key)}
    end
  end

  def to_array(prefix = '')
    array = []
    each{|key| array << prefix + key}
    array
  end

  def inspect
    "#<Trie: {#{to_array.join(', ')}}>"
  end

  protected
    def self.shift(key)
      [key[0], key[1..-1]]
    end

    def self.join(value, key)
      key.is_a?(String) and value.respond_to?(:chr) ?
        value.chr + key : key.unshift(value)
    end

    def visit_node(key, &block)
      car, cdr = Trie.shift(key) unless key.empty?
      if child = @children[car]
        child.visit_node(cdr, &block)
      else
        block.call(self, key)
      end
    end
    
    def append(key)
      car, cdr = Trie.shift(key)
      if car
        @children[car] = Trie.new
        cdr.empty? ? @children[car].terminal! : @children[car].append(cdr)
      end
    end
    
    def prune(value)
      child = @children[value]
      if child.terminal?
        child.leaf? ? @children.delete(value) : child.internal!
        child
      end
    end

    def terminal!
      @terminal = true
    end
    
    def internal!
      @terminal = false
    end
    
    def terminal?
      @terminal
    end
    
    def leaf?
      !@children.any?
    end
end

class WordFinder
  Prefix = Struct.new(:value, :row, :column)

  def initialize(input_file, words_file)
    lines = nil
    File.open(input_file) { |f| lines = f.read.each_line }

    # first line in the file contains the size of table
    size = lines.next.to_i
    @max = size - 1

    @table = size.times.each_with_object([]) do |row, result|
      result << lines.next.chomp.chars.to_a.each_with_index.map do |char, column|
        Prefix.new(char, row, column)
      end
    end

    # make prefixes with two letters
    @prefixes = @table.flatten.each_with_object([]) do |m, result|
      get_surrounds(m).each do |s|
        result << Prefix.new(m.value+s.value, s.row, s.column)
      end
    end

    @trie = Trie.new
    File.open(words_file) { |f| f.read.each_line { |line| @trie << line.chomp } }
  end

  def generate
    f = File.open('output.txt', 'w')

    at_exit { f.flush; f.close }

    until @prefixes.empty?
      @prefixes = @prefixes.each_with_object([]) do |m, result|

        get_surrounds(m).each do |prefix|

          word = m.value + prefix.value     

          unless @trie.prefixed_by(word).empty?
            result << Prefix.new(word, prefix.row, prefix.column)
            if @trie.include?(word)
              f.puts word 
              @trie.delete(word)
            end
          end           
        end
      end 
    end
  end

  private

  def get_surrounds(prefix)
    row    = prefix.row
    column = prefix.column

    prefixes = []
    prefixes << @table[row-1][column]   if row > 0         # upper element
    prefixes << @table[row+1][column]   if row < @max      # down element
    prefixes << @table[row][column-1]   if column > 0      # left element      
    prefixes << @table[row][column+1]   if column < @max   # right element
    prefixes << @table[row+1][column+1] if row < @max and column < @max   # rd
    prefixes << @table[row-1][column+1] if row > 0 and column < @max      # ru
    prefixes << @table[row+1][column-1] if row < @max and column > 0      # ld
    prefixes << @table[row-1][column-1] if row > 0 and column > 0         # lu

    prefixes      
  end
end

WordFinder.new('input.txt', 'words.txt').generate