Skip to content

Instantly share code, notes, and snippets.

@JoelQ

JoelQ/tree_enumerator.rb

Created August 21, 2022 22:15
Show Gist options
  • Save JoelQ/02f3ef9f61bebc7c8e5ea67d10ed92c6 to your computer and use it in GitHub Desktop.
Save JoelQ/02f3ef9f61bebc7c8e5ea67d10ed92c6 to your computer and use it in GitHub Desktop.
Download ZIP
Demonstration of how using `Enumerator` makes it nicer to work with a data structure that has multiple valid traversals such as a binary tree
Raw
tree_enumerator.rb
class Tree
attr_reader :val, :left, :right
def self.empty
EmptyTree.new
end
def self.leaf(val)
new(val, empty, empty)
end
def initialize(val, left, right)
@val = val
@left = left
@right = right
end
# ==== Traversals ====
#
# These methods provide different ways of traversing a tree. See
# https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search for more on
# the different ways a tree can be traversed.
#
# The methods can be called with a block and each item will be yielded in that
# given order:
#
# tree.preorder { |n| puts n }
#
# When called without a block they return a `Enumerator`. This allows calling
# any of the `Enumerable` methods for that given traversal such as:
#
# tree.inorder.map { |n| n * 2 }
# tree.postorder.find(&:even?)
# Depth-first search, pre-order. Given this tree:
#
# 1
# / \
# 2 5
# / \ / \
# 3 4 6 7
#
# A pre-order traversal will go:
# 1, 2, 3, 4, 5, 6, 7
def preorder(&block)
return to_enum(:preorder) unless block_given?
block.call(val)
left.preorder(&block)
right.preorder(&block)
end
# Depth-first search, in-order. Given this tree:
#
# 1
# / \
# 2 5
# / \ / \
# 3 4 6 7
#
# An in-order traversal will go:
# 3, 2, 4, 1, 6, 5, 7
def inorder(&block)
return to_enum(:inorder) unless block_given?
left.inorder(&block)
block.call(val)
right.inorder(&block)
end
# Depth-first search, post-order. Given this tree:
#
# 1
# / \
# 2 5
# / \ / \
# 3 4 6 7
#
# A post-order traversal will go:
# 3, 4, 2, 6, 7, 5, 1
def postorder(&block)
return to_enum(:postorder) unless block_given?
left.postorder(&block)
right.postorder(&block)
block.call(val)
end
end
class EmptyTree
def preorder
end
def postorder
end
def inorder
end
end
# 1
# / \
# 2 5
# / \ / \
# 3 4 6 7
tree =
Tree.new(
1,
Tree.new(
2,
Tree.leaf(3),
Tree.leaf(4),
),
Tree.new(
5,
Tree.leaf(6),
Tree.leaf(7),
),
)
# (2 * 3) + (12 / 6)
expression =
Tree.new(
:+,
Tree.new(:*, Tree.leaf(2), Tree.leaf(3)),
Tree.new(:/, Tree.leaf(12), Tree.leaf(6)),
)
@camertron
Copy link

This is awesome, thanks for sharing! A humble suggestion: I often use the __method__ local when turning a method into an enum, eg:

def postorder(&block)
  return to_enum(__method__) unless block_given?
  ...
end

Renaming the method is easier now, since I only have to change it in one place. It's especially handy if you're calling to_enum farther down in the method body so it's less visually obvious that you have to change the method name in multiple places. I've been bitten by this in the past so I thought I'd mention it 😄

@JoelQ
Copy link
Author

JoelQ commented Aug 15, 2024

@camertron that's a great tip! Thanks for sharing!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment