lisovskyvlad · January 8, 2017 18:15 · Jan 8, 2017
diff --git a/tries.rb b/tries.rb
@@ -0,0 +1,74 @@
+class Trie
+  Node = Struct.new(:char, :children, :is_complete_word)
+
+  attr_reader :root
+
+  def initialize
+    @root = Node.new('', {}, false)
+  end
+
+  def add_word(word)
+    char_array = word.split('')
+    add_char(root, char_array)
+  end
+
+  def count_prefixes(word)
+    char_array = word.split('')
+    words_with_prefix(root, char_array)
+  end
+
+  private
+
+  def add_char(node, char_array)
+    char = char_array.shift
+    if char
+      child_node = node.children[char]
+
+      if child_node
+        add_char(child_node, char_array)
+      else
+        new_child_node = Node.new(char, {}, false)
+        node.children[char] = new_child_node
+        add_char(new_child_node, char_array)
+      end
+    else
+      node.is_complete_word = true
+    end
+  end
+
+  def words_with_prefix(node, char_array)
+    char = char_array.shift
+    if char
+      child_node = node.children[char]
+      if child_node
+        words_with_prefix(child_node, char_array)
+      else
+        0
+      end
+    else # prefix is ended
+      arr = []
+      count_completed_words(node, arr)
+      arr.length
+    end
+  end
+
+  def count_completed_words(node, out_arr)
+    node.children.each do |char, child_node|
+      count_completed_words(child_node, out_arr)
+    end
+    out_arr.push(node) if node.is_complete_word
+  end
+end
+
+trie = Trie.new
+
+n = gets.strip.to_i
+for a0 in (0..n-1)
+  op, contact = gets.strip.split(' ')
+  case op
+  when 'add'
+    trie.add_word(contact)
+  when 'find'
+    puts trie.count_prefixes(contact)
+  end
+end