T9 Kata

t9.rb

##################################################################
# Given a list of words, model a T9[1] predictive text engine.   #
#                                                                #
# Input to the function will be a series of pressed digits. The  #
# output should be a set of all possible words that can be built #
# using those keys as the prefix.                                #
#                                                                #
# [1] http://en.wikipedia.org/wiki/T9_(predictive_text)          #
##################################################################

require 'set'

# TODO: extract out the Trie/prefix tree as a separate concept.
class T9
  KEYPAD = {}

  [
    %w[. , ? ! + -], %w[a b c], %w[d e f],
    %w[g h i],       %w[j k l], %w[m n o],
    %w[p q r s],     %w[t u v], %w[w x y z]
  ].each_with_index do |letters, index|
    letters.each do |letter|
      key = index + 1
      KEYPAD[letter] = key
    end
  end

  attr_accessor :children, :words

  def initialize
    @children = {}
    @words    = Set.new
  end
  
  def insert(word, path = word.downcase.split(//))
    if path.empty?
      words << word
    else 
      first, *rest = path
      digit        = KEYPAD[first]
      child        = children[digit] ||= self.class.new

      child.insert word, rest
    end

    word
  end
  
  def descendants
    children = self.children.values.compact
    children + children.flat_map(&:descendants)
  end

  def all_words
    words + descendants.map(&:words).reduce(&:union)
  end

  def dial(digits = [])
    tree      = children[digits.first]
    remaining = digits.drop(1)

    return Set.new unless tree
    
    if remaining.empty?
      tree.all_words
    else  
      tree.dial(remaining)
    end
  end
end

if $0 == __FILE__ 
  require 'minitest/autorun'

  describe T9 do
    let(:words) { %w[hello he id identity apple] }
    let(:t9)    { T9.new }

    before do
      words.each { |word| t9.insert(word) }
    end

    it { t9.dial([]).must_equal     Set.new }
    it { t9.dial([7]).must_equal    Set.new }
    it { t9.dial([2]).must_equal    Set.new(%w[apple]) }
    it { t9.dial([4, 3]).must_equal Set.new(%w[he hello id identity]) }
  end
end