Trie w/ Ruby

trie.rb

#!/bin/ruby

class Trie
    @@root

    def initialize
        @@root = TrieNode.new
    end

    def add_word(word)
        current = @@root

        word.split('').each do |chr|
            node = current.get_node(chr)

            if !node.is_a?(TrieNode)
                node = TrieNode.new
                current.add_node(chr, node)
            end

            current = node
        end

        current.set_end_of_word
    end

    def count_prefix_words(prefix)
        current = @@root

        prefix.split('').each do |chr|
            node = current.get_node(chr)

            return 0 if !node.is_a?(TrieNode)

            current = node
        end

        children = current.get_children

        return walk_and_count_word_ends(children)
    end

    def walk_and_count_word_ends(array)
        count = 0
class Trie
    @@root

    def initialize
        @@root = TrieNode.new
    end

    def add_word(word)
        current = @@root

        word.split('').each do |chr|
            node = current.get_node(chr)

            if !node.is_a?(TrieNode)
                node = TrieNode.new
                current.add_node(chr, node)
            end

            current = node
        end

        current.set_end_of_word
    end

    def count_prefix_words(prefix)
        current = @@root

        prefix.split('').each do |chr|
            node = current.get_node(chr)

            return 0 if !node.is_a?(TrieNode)

            current = node
        end

        children = current.get_children

        return walk_and_count_word_ends(children)
    end

    def walk_and_count_word_ends(array)
        count = 0

        array.each do |key, item|
            if item.is_end_of_word
                count += 1
            end
            
            count += walk_and_count_word_ends(item.get_children)
        end

        return count
    end

    def output
        puts @@root.get_children
    end
end

class TrieNode

    def initialize
        @children = Hash.new
        @end_of_word = false    
    end

    def get_node(chr)
        return @children[chr] if @children.key?(chr)
    end

    def add_node(chr, node)
        @children[chr] = node
    end

    def is_end_of_word
        @end_of_word
    end

    def set_end_of_word
        @end_of_word = true
    end

    def get_children
        @children
    end

    def has_children
        @children.length > 0
    end

    attr_accessor :children
    attr_accessor :end_of_word
end

trie = Trie.new()
trie.add_word("happy")
trie.add_word("ham")
trie.add_word("hart")

puts trie.count_prefix_words("h") # 3