splay_tree.cpp

#ifndef SPLAY_TREE_H
#define SPLAY_TREE_H

#include <utility> 
#include <iostream>
#include "Array.h"
#include "bst.h"
#include <cassert>

using namespace std;

template<typename KeyT, typename ValueT>
class SplayTree : public BST<KeyT,ValueT> {
  public:  
    enum Orientation {
        NONE, LEFT, RIGHT
    };
    /// Find a node in the BST containing the given key
    BSTNode<KeyT,ValueT>* find(KeyT aKey)
    {
        BSTNode<KeyT,ValueT>* node = BST<KeyT,ValueT>::find(aKey);
        if (node != nullptr) {
            cout << "doing a splay " << endl;
            // Splay the node that is found or the node where the search ended
            splay(node);
        }
        return node;
    }

    /// Find a node in the BST containing the given key.
    /// If exact match not found, return nullptr but also set the approxMatchNode
    BSTNode<KeyT,ValueT>* findApprox(KeyT aKey, 
                                     BSTNode<KeyT,ValueT>*& approxMatchNode_lb,
                                     BSTNode<KeyT,ValueT>*& approxMatchNode_ub)
    {
        BSTNode<KeyT,ValueT>* node = BST<KeyT,ValueT>::findApprox(aKey, 
                                                                  approxMatchNode_lb,
                                                                  approxMatchNode_ub);
        if (node != nullptr) {
            // Splay the node that is found or the node where the search ended
            splay(node);
        } else {
	  if (approxMatchNode_lb)
            splay(approxMatchNode_lb);
	  if (approxMatchNode_ub)
	    splay(approxMatchNode_ub);
        }
        return node;
    }
  
    /// Find min
    BSTNode<KeyT,ValueT>* getMin() 
    { 
        BSTNode<KeyT,ValueT>* minBSTNode = BST<KeyT,ValueT>::getMin(); 
        if (minBSTNode != nullptr) {
            // Splay the min node
            splay(minBSTNode);
        }
        return minBSTNode;
    }

    /// Find max
    BSTNode<KeyT,ValueT>* getMax() 
    { 
        BSTNode<KeyT,ValueT>* maxBSTNode = BST<KeyT,ValueT>::getMax(); 
        if (maxBSTNode != nullptr) {
            // Splay the max node
            splay(maxBSTNode);
        }
        return maxBSTNode;
    }

    /// Remvoe
    void remove(KeyT aKey) 
    { 
        BSTNode<KeyT,ValueT>* parentNode = BST<KeyT,ValueT>::remove(aKey); 
	if (parentNode != nullptr) {
            // Splay the parent node of the removed node
            splay(parentNode);
        }
    }


  public:

    // 1. if aNode < root (zig), get left child or vice versa
    // 2. set the parent of old root as aNode
    // 3. if aNode < root (zig) set Left child as previous
    BSTNode<KeyT,ValueT>* setRotationValues(bool zig, BSTNode<KeyT,ValueT>* aNode, BSTNode<KeyT,ValueT>* aParent) {
       BSTNode<KeyT,ValueT>* aChild = zig ? aNode->left() : aNode->right();
       zig ? aNode->setRight(aParent) : aNode->setLeft(aParent);
       aParent->setParent(aNode);
       zig ? aParent->setLeft(aChild) : aParent->setRight(aChild);
       return aChild;
    }

    void updateRoot(BSTNode<KeyT,ValueT>* aNode, BSTNode<KeyT,ValueT>* aParent, BSTNode<KeyT,ValueT>* aChild) {
        if (aChild != nullptr) aChild->setParent(aParent);
        if (this->root() == aParent) {
            this->setRoot(aNode); // Set the new root
            aNode->setParent(nullptr);
        }
    }

    // Zig rotation -- left to right. Use-case: aNode is a left child of root
    void zig(BSTNode<KeyT,ValueT>* aNode, BSTNode<KeyT,ValueT>* aParent) {
        BSTNode<KeyT,ValueT>* rightChild = setRotationValues(true, aNode, aParent);
        updateRoot(aNode, aParent, rightChild);
    }
    
    // Zag rotation -- right to left. Use-case: aNode is a right child of root
    void zag(BSTNode<KeyT,ValueT>* aNode, BSTNode<KeyT,ValueT>* aParent) {
        BSTNode<KeyT,ValueT>* leftChild = setRotationValues(false, aNode, aParent);
        updateRoot(aNode, aParent, leftChild);
    }

    void doZigOrZag(BSTNode<KeyT,ValueT>* aNode, BSTNode<KeyT,ValueT>* aParent) {
      isZig(aNode) ? zig(aNode, aParent) : zag(aNode, aParent);
    }

    Orientation getOrientation(BSTNode<KeyT,ValueT>* gParent, BSTNode<KeyT,ValueT>* ggParent) {
      if (ggParent != nullptr) { // if grandparent is not root
        return (ggParent->left() == gParent ? LEFT : RIGHT); 
      }
      return NONE;
    }

    bool isZig(BSTNode<KeyT,ValueT>* aNode) { return aNode->parent()->left() == aNode; }

    pair<int,int> checkDouble(BSTNode<KeyT,ValueT>* aNode) {
      pair<int,int> directions;
      directions.first = isZig(aNode); // check child node orientation
      directions.second = isZig(aNode->parent()); // check parent node orientation
      return directions;
    }

    // Double rotation -- zigzig/zagzag/zagzig/zigzag. Use-case: aNode is a grandchild
    void doubleZigOrZag(BSTNode<KeyT,ValueT>* aNode, BSTNode<KeyT,ValueT>* aParent, BSTNode<KeyT,ValueT>* aGrandParent) {
      pair<int,int> directions = checkDouble(aNode); // note the current state
      BSTNode<KeyT,ValueT>* ggParent = aGrandParent->parent(); // get great grandparent
      Orientation orient = getOrientation(aGrandParent, ggParent);
      directions.first == 1 ? zig(aNode, aParent) : zag(aNode, aParent); // Rotate around parent
      directions.second == 1 ? zig(aNode, aParent) : zag(aNode, aParent); // Rotate around grandparent
      if (orient != NONE) {
        orient == LEFT ? ggParent->setLeft(aNode) : ggParent->setRight(aNode);
        aNode->setParent(ggParent); // aNode is now root. Set ggParent as parent.
      } else {
        aNode->setParent(nullptr); 
      }
    }

  // Splay operation
  void splay(BSTNode<KeyT,ValueT>* aNode) {
    BSTNode<KeyT,ValueT>* parent = aNode->parent();
    if (!parent) return; // already root
    BSTNode<KeyT,ValueT>* grandParent = parent->parent();
    while (parent != nullptr && grandParent != nullptr) { 
      doubleZigOrZag(aNode, parent, grandParent); 
      parent = aNode->parent(); // Evaluate new parent and grandparent after rotations
      if (parent == nullptr) break;
      grandParent = parent->parent(); 
      if (grandParent == nullptr) break;
    }
    // Zig operation is done only as a last pass
    // after Zigzig and/or Zigzag operations
    if (grandParent == nullptr) doZigOrZag(aNode, parent); // Parent is the root -- Zig operation
  }

  // TODO: Generate a value for each string and make a comparison that way.
  // Otherwise make a custom comparator operator function
  BSTNode<KeyT,ValueT>* insert(int aKey) {
    BSTNode<KeyT,ValueT>* node = this->insert(aKey);
    if (node != nullptr) {
      // Splay the node that is found or the node where the search ended
      splay(node);
    }
    return node;
  }

  void splay2(BSTNode<KeyT,ValueT>* aNode) 
  {
    BSTNode<KeyT,ValueT>* parent = aNode->parent(); // 1. get parent
    if (!parent) return;
    BSTNode<KeyT,ValueT>* grandParent = parent->parent(); // 2. get grandparent.
    while (parent != nullptr && grandParent != nullptr) { // 3. while both exist,
      doubleZigOrZag(aNode, parent, grandParent); // 3a. we know at least double is necessary

    }
  }
};


#endif