binary_search_tree.rb

require 'rspec'

module BST
  class BSTNode
    attr_accessor :key, :value, :left, :right

    def copy(other)
      @key = other.key
      @value = other.value
      @left = other.left
      @right = other.right

      self
    end

    def copy_key_and_value(other)
      @key = other.key
      @value = other.value
      self
    end

    def set(key, value)
      @key, @value = key, value
      add_children unless has_children?
      value
    end

    def has_left_child?
      not @left.key.nil?
    end

    def find_right_most_descendant
      return self if @right.key.nil?
      @right.find_right_most_descendant
    end

    private

    def add_children
      @left = BSTNode.new
      @right = BSTNode.new
    end

    def has_children?
      not (@left.nil? || @right.nil?)
    end
  end

  class SearchTree
    def initialize
      @root_node = BSTNode.new
    end

    def [](key)
      search(key).value
    end

    def search(key)
      search_at(@root_node, key)
    end

    def []=(key, value)
      return unset(key) if value.nil?
      search_at(@root_node, key).set(key,value)
    end

    def unset(key)
      node = search_at(@root_node, key)
      if node.has_left_child?
        swap_with_previous_node_and_remove(node)
      else
        node.copy(node.right)
      end

      nil
    end

    def traverse(&block)
      traverse_at(@root_node, &block)
    end

    def each
      traverse do |node|
        yield node.key, node.value if block_given?
      end
    end

    private

    def search_at(search_node, key)
      return search_node if search_node.key.nil? || search_node.key == key

      if search_node.key < key
        search_at(search_node.right, key)
      else
        search_at(search_node.left, key)
      end
    end

    def traverse_at(node, &block)
      return if node.key.nil?

      traverse_at(node.left, &block)
      yield node if block_given?
      traverse_at(node.right, &block)
    end

    def swap_with_previous_node_and_remove(node)
      previous_node = node.left.find_right_most_descendant

      node.copy_key_and_value(previous_node)
      previous_node.copy(previous_node.left)
    end

  end
end

describe BST::SearchTree do
  describe '#[]' do
    it 'should return nil when the tree is empty' do
      expect(subject['key']).to be_nil
    end

    it 'should return the root nodes value when we look it up' do
      subject['key'] = 'value'
      expect(subject['key']).to eq('value')
    end

    it 'should work with 3 values inserted in ascending order' do
      subject[1] = 'value1'
      subject[2] = 'value2'
      subject[3] = 'value3'
      expect(subject[1]).to eq('value1')
      expect(subject[2]).to eq('value2')
      expect(subject[3]).to eq('value3')
    end

    it 'should work with 3 values inserted in descending order' do
      subject[3] = 'value3'
      subject[2] = 'value2'
      subject[1] = 'value1'
      expect(subject[1]).to eq('value1')
      expect(subject[2]).to eq('value2')
      expect(subject[3]).to eq('value3')
    end

    it 'should overwrite previous value' do
      subject[4] = 'old_value'
      subject[4] = 'new_value'
      expect(subject[4]).to eq('new_value')
    end
  end

  describe '#search' do
    it 'should return the actual node' do
      subject[3] = 'value2'
      subject[1] = 'value1'
      subject[0] = 'value0'
      subject[2] = 'value1.5'

      node = subject.search(1)

      expect(node.key).to eq(1)
    end
  end

  describe '#traverse' do
    it 'should yield all elements in order' do
      subject[2] = 'value2'
      subject[1] = 'value1'
      subject[0] = 'value0'
      subject[3] = 'value3'

      result = []
      subject.traverse do |node|
        result << "#{node.key} = #{node.value}"
      end
      expect(result.join ', ').to eq('0 = value0, 1 = value1, 2 = value2, 3 = value3')
    end
  end

  describe '#each' do
    it 'should yield all elements in order' do
      subject[2] = 'value2'
      subject[1] = 'value1'
      subject[0] = 'value0'
      subject[3] = 'value3'

      result = []
      subject.each do |key, value|
        result << "#{key} = #{value}"
      end
      expect(result.join ', ').to eq('0 = value0, 1 = value1, 2 = value2, 3 = value3')
    end
  end

  describe '#unset' do
    it 'should remove the single element' do
      subject[1] = 'value1'
      subject.unset(1)
      expect(subject[1]).to be_nil
    end

    it 'should remove a leaf node, but it should not remove the root' do
      subject[1] = 'value1'
      subject[2] = 'value2'
      subject.unset(2)
      expect(subject[1]).to eq('value1')
      expect(subject[2]).to be_nil
    end

    it 'should remove non leaf element' do
      subject[2] = 'value2'
      subject[1] = 'value1'
      subject[0] = 'value0'
      subject[1.5] = 'value1.5'
      subject[3] = 'value3'

      subject.unset(2)

      expect(subject[0]).to eq('value0')
      expect(subject[1]).to eq('value1')
      expect(subject[1.5]).to eq('value1.5')
      expect(subject[2]).to be_nil
      expect(subject[3]).to eq('value3')

    end
  end
end