ruby 在Ruby中实现Trie

Posted 2021-05-15

# A basic ruby implementation for the trie structure.
class Trie
  # Node class to store node information (character or tag, and action)
  class Node
    attr_reader :children, :tag
    attr_accessor :value

    def initialize(tag, parent, action = nil)
      @tag = tag
      parent.children << self if parent
      @children = []
      @action = action
    end

    def find_child(tag)
      children.find { |node| node.tag == tag }
    end
  end

  # Initialize the trie with an empty root node
  def initialize(name = nil)
    @root = Node.new(name, nil)
  end

  # Add a new set of nodes to the trie.
  def add(string, value)
    current = @root

    string.chars.each do |tag|
      current = current.find_child(tag) || Node.new(tag, current, nil)
    end

    current.value = value
  end

  # Execute a node action given a string. If no node is found, nothing is executed.
  def value(string)
    node = find_node(string)
    node.value if node
  end

  # Pretty print the trie recursively
  def pretty_print(current = @root, level = 0, indent = ' ')
    puts "#{indent * level}#{current.tag}#{current.value ? " (#{current.value})" : ''}"
    current.children.each { |child| pretty_print(child, level + 1, indent) }
  end

  private

  # Traverse the trie to find a node
  def find_node(string)
    current = @root

    string.chars.each do |tag|
      current = current.find_child(tag)
      return if current.nil?
    end

    current
  end
end

# Run example
# ---

trie = Trie.new 'trie_example'

trie.add 'zero', 0
trie.add 'one', 1
trie.add 'two', 2
trie.add 'point', ->(x, y) { "Point #{x},#{y}" }

puts trie.value('zero'), trie.value('one'), trie.value('two'), trie.value('three')
puts trie.value('point').call(41, 42)

trie.pretty_print