binary_search_tree.rb

require "rspec/autorun"

# An implementation of an immutable binary search tree in Ruby, with a small
# suite of tests. Supports #==, #insert, #delete, and #to_a (an in-order
# traversal).
#
# To run the tests, just evaluate the file: `ruby binary_search_tree.rb`

class Leaf
  def ==(other)
    self.class == other.class
  end

  def to_a
    []
  end

  def leaf?
    true
  end

  def insert(value)
    BinarySearchTree.new(value: value)
  end
end

class BinarySearchTree
  def initialize(value:, left: Leaf.new, right: Leaf.new)
    @value = value
    @left = left
    @right = right
  end

  def ==(other)
    self.class == other.class &&
      value == other.value &&
      left == other.left &&
      right == other.right
  end

  def insert(new_value)
    if new_value <= value
      self.class.new(
        value: value,
        left: left.insert(new_value),
        right: right,
      )
    elsif new_value > value
      self.class.new(
        value: value,
        left: left,
        right: right.insert(new_value),
      )
    end
  end

  def delete(target)
    if target < value
      self.class.new(
        value: value,
        left: left.delete(target),
        right: right,
      )
    elsif target > value
      self.class.new(
        value: value,
        left: left,
        right: right.delete(target),
      )
    else
      delete_self
    end
  end

  def to_a
    left.to_a + [value] + right.to_a
  end

  def leaf?
    false
  end

  def leftmost_value
    if left.leaf?
      value
    else
      left.leftmost_value
    end
  end

  protected

  attr_reader :value, :left, :right

  private

  def delete_self
    if two_children?
      self.class.new(
        value: right.leftmost_value,
        left: left,
        right: right.delete(right.leftmost_value),
      )
    elsif left.leaf?
      right
    else
      left
    end
  end

  def two_children?
    !left.leaf? && !right.leaf?
  end
end

describe BinarySearchTree do
  def sample_tree
    #   5
    #  / \
    # 3   6
    #      \
    #       8
    BinarySearchTree.new(
      value: 5,
      left: BinarySearchTree.new(
        value: 3,
      ),
      right: BinarySearchTree.new(
        value: 6,
        right: BinarySearchTree.new(
          value: 8,
        ),
      ),
    )
  end

  describe "#to_a" do
    it "performs an in-order traversal" do
      #   5
      #  / \
      # 3   6    => [3, 5, 6, 8]
      #      \
      #       8
      expect(sample_tree.to_a).to eq([3, 5, 6, 8])
    end
  end

  describe "#==" do
    it "returns true if two trees are the same" do
      tree_a = sample_tree
      tree_b = sample_tree

      expect(tree_a).to eq(tree_b)
    end
  end

  describe "#insert" do
    it "returns a tree with the new value" do
      #   5            5
      #  / \     +7   / \
      # 3   6    =>  3   6
      #      \            \
      #       8            8
      #                   /
      #                  7
      expected = BinarySearchTree.new(
        value: 5,
        left: BinarySearchTree.new(
          value: 3,
        ),
        right: BinarySearchTree.new(
          value: 6,
          right: BinarySearchTree.new(
            value: 8,
            left: BinarySearchTree.new(
              value: 7,
            ),
          ),
        ),
      )

      expect(sample_tree.insert(7)).to eq(expected)
    end
  end

  describe "#delete" do
    it "returns a tree without the old value if the target node has no children" do
      #   5            5
      #  / \     -8   / \
      # 3   6    =>  3   6
      #      \
      #       8
      expected = BinarySearchTree.new(
        value: 5,
        left: BinarySearchTree.new(
          value: 3,
        ),
        right: BinarySearchTree.new(
          value: 6,
        ),
      )

      expect(sample_tree.delete(8)).to eq(expected)
    end

    it "returns a tree without the old value if the target node has one child" do
      #   5            5
      #  / \     -6   / \
      # 3   6    =>  3   8
      #      \
      #       8
      expected = BinarySearchTree.new(
        value: 5,
        left: BinarySearchTree.new(
          value: 3,
        ),
        right: BinarySearchTree.new(
          value: 8,
        ),
      )

      expect(sample_tree.delete(6)).to eq(expected)
    end

    it "returns a tree without the old value if the target node has two children" do
      #   5            6
      #  / \     -5   / \
      # 3   6    =>  3   8
      #      \
      #       8
      expected = BinarySearchTree.new(
        value: 6,
        left: BinarySearchTree.new(
          value: 3,
        ),
        right: BinarySearchTree.new(
          value: 8,
        ),
      )

      expect(sample_tree.delete(5)).to eq(expected)
    end
  end
end