Binary Search Tree Java Implementation

BinarySearchTree.java

import java.util.ArrayList;
import java.util.stream.Collectors;

class Node<K extends Comparable<K>, V> {
    K key;
    V value;
    Node<K, V> parent;
    Node<K, V> left;
    Node<K, V> right;

    public Node(K key, V value) {
        this.key = key;
        this.value = value;
    }

    public Node(K key, V value, Node<K, V> parent) {
        this(key, value);
        this.parent = parent;
    }
    
    public void inOrderTraverse() {
        if (left != null)
            left.inOrderTraverse();
        System.out.println(this);
        if (right != null)
            right.inOrderTraverse();
    }

    public void preOrderTraverse() {
        System.out.println(this);
        if (left != null)
            left.preOrderTraverse();
        if (right != null)
            right.preOrderTraverse();
    }

    @Override
    public String toString() {
        return String.format("(%s, %s)", this.key, this.value);
    }

    public int compareTo(Node<K, V> other) {
        return this.key.compareTo(other.key);
    }

    public boolean hasBothChildren() {
        return hasLeftChild() && hasRightChild();
    }

    public boolean hasLeftChild() {
        return this.left != null;
    }

    public boolean hasRightChild() {
        return this.right != null;
    }

    public boolean isLeaf() {
        return !hasLeftChild() && !hasRightChild();
    }

}

public class BinarySearchTree<K extends Comparable<K>, V> {
    private Node<K, V> root;

    private boolean debug = false;

    /**
     * Method takes a key and a value as input and inserts the key/value pair into
     * its correct position in the tree.
     *
     * @param key
     * @param value
     */
    public void add(K key, V value) {
        Node<K, V> toAdd = new Node<K, V>(key, value);
        if (debug) {
            System.out.println("adding '" + key + "'");
        }
        if (this.root == null) {
            this.root = toAdd;
            return;
        }
        addHelper(this.root, toAdd);
    }

    /**
     * Helper method for add; Recursively searches for appropriate location to
     * insert the toAdd node.
     *
     * @param root
     * @param toAdd
     */
    private void addHelper(Node<K, V> root, Node<K, V> toAdd) {
        int result = toAdd.compareTo(root);
        if (result == 0) {
            // key is already in tree, so existing value is replaced.
            if (debug) {
                System.out.println("Replacing existing value \'" + root.value + "\' with \'" + toAdd.value + "\'");
            }
            root.value = toAdd.value;
        } else {
            // set as parent
            toAdd.parent = root;
            if (result > 0) {
                if (root.right == null) {
                    root.right = toAdd;
                } else {
                    addHelper(root.right, toAdd);
                }
            } else if (result < 0) {
                if (root.left == null) {
                    root.left = toAdd;
                } else {
                    addHelper(root.left, toAdd);
                }
            }
        }
    }

    /**
     * Method removes and returns the value associated with the given key; return
     * null if the key is not found.
     *
     * @param key
     * @return
     */
    public V remove(K key) {
        if (debug) {
            System.out.println("removing '" + key + "'.");
        }
        Node<K, V> toRemove = search(key);

        // Case: Key is not found.
        if (toRemove == null) {
            // key was not found.
            System.out.println("key not found.");
            return null;
        }
        V toReturn = toRemove.value;

        // Case: toRemove is a leaf node.
        if (toRemove.isLeaf()) {
            Node<K, V> parent = toRemove.parent;
            if (debug) {
                System.out.println("removing a leaf node...");
                System.out.println("parent = " + parent);
            }
            if (parent != null) {
                if (parent.left == toRemove) {
                    parent.left = null;
                } else if (parent.right == toRemove) {
                    parent.right = null;
                }
            } else {
                if (debug) {
                    System.out.println("The root has null parent");
                }
                if (isRoot(toRemove)) {
                    this.root = null;
                }
            }

            return toReturn;
        }

        // Case: toRemove has both children
        if (toRemove.hasBothChildren()) {

            // Find the successor node
            Node<K, V> successor = toRemove.right;

            boolean hasAdvancedSuccessorReference = false;

            if (debug) {
                System.out.println("Finding successor node...");
            }

            while (successor != null && successor.left != null) {
                successor = successor.left;
                // If true, then the left-subtree of the successor has been explored and the
                // smallest value has been found.
                hasAdvancedSuccessorReference = true;
            }

            // Replace toRemove information with the successor's information
            toRemove.key = successor.key;
            toRemove.value = successor.value;

            if (hasAdvancedSuccessorReference) {
                System.out.println("Removing node with both children...");
                // Remove duplicate
                successor.parent.left = successor.right;
            } else {
                // The successor does not have a left sub-branch, so the successor itself is the
                // smallest value available.
                successor.parent.right = successor.right;
            }

            return toReturn;
        }

        // Case: toRemove has a left child
        if (toRemove.hasLeftChild()) {
            if (debug) {
                System.out.println("Removing node with left child...");
            }
            Node<K, V> parent = toRemove.parent;
            // Corner case: toRemove is the root
            if (isRoot(toRemove)) {
                if (debug) {
                    System.out.println("Removing the root...");
                }
                this.root = this.root.left;
                // Update parent reference
                this.root.parent = null;
            } else {
                if (parent.left == toRemove) {
                    parent.left = toRemove.left;
                    // Update parent reference
                    parent.left.parent = parent;
                } else if (parent.right == toRemove) {
                    parent.right = toRemove.left;
                    // Update parent reference
                    parent.right.parent = parent;
                }
            }

            return toReturn;
        }

        // Case: toRemove has a right child
        if (toRemove.hasRightChild()) {
            if (debug) {
                System.out.println("Removing node with right child...");
            }
            Node<K, V> parent = toRemove.parent;
            // Corner case: toRemove is the root
            if (isRoot(toRemove)) {
                if (debug) {
                    System.out.println("Removing the root...");
                }
                this.root = this.root.right;
                // Update parent reference
                this.root.parent = null;
            } else {
                if (parent.left == toRemove) {
                    parent.left = toRemove.right;
                    // Update parent reference
                    parent.left.parent = parent;
                } else if (parent.right == toRemove) {
                    parent.right = toRemove.right;
                    // Update parent reference
                    parent.right.parent = parent;
                }
            }
            return toReturn;
        }

        return null;
    }

    /**
     * Method searches for the node that contains the given key and returns the node
     * if found; otherwise, returns null.
     *
     * @param key
     * @return
     */
    private Node<K, V> search(K key) {
        return searchHelper(this.root, key);
    }

    /**
     * Helper method for search; Recursively searches for the node that contains the
     * given key and returns the node if found; otherwise, returns null.
     *
     * @param root
     * @param key
     * @return
     */
    private Node<K, V> searchHelper(Node<K, V> root, K key) {
        if (root == null) {
            if (debug) {
                System.out.println("key not found.");
            }
            return null;
        }

        int result = key.compareTo(root.key);
        if (result == 0) {
            return root;
        } else if (result > 0) {
            return searchHelper(root.right, key);
        } else if (result < 0) {
            return searchHelper(root.left, key);
        }
        if (debug) {
            System.out.println("key not found.");
        }
        return null;
    }

    /**
     * Method determines if the given key is present or not in the tree; returns
     * associated value if given key is present, otherwise returns null.
     *
     * @param key
     * @return
     */
    public V lookup(K key) {
        Node<K, V> node = search(key);
        return node == null ? null : node.value;
    }

    /**
     * Traverses the tree starting from the root and prints all key/value pairs in
     * increasing order.
     */
    public void inOrderTraverse() {
        if (this.root != null)
            this.root.inOrderTraverse();
    }

    /**
     * Traverses the tree starting from the root and prints all key/value pairs in
     * the following order: 1. Prints parent 2. Prints left 3. Prints right
     */
    public void preOrderTraverse() {
        if (this.root != null)
            this.root.preOrderTraverse();

    }

    /**
     * NOTE: This method was added to improve readability. Determines if the given
     * node is the root.
     *
     * @param n
     * @return
     */
    private boolean isRoot(Node<K, V> n) {
        return this.root == n;
    }

    @Override
    public String toString() {
        if(this.root == null) return "";
        ArrayList<Node<K, V>> nodes = new ArrayList<>();
      
        toStringHelper(this.root, nodes);
        String representationString = "[";
      
        representationString += nodes.stream()
          .map(aNode -> aNode.toString())
          .collect(Collectors.joining(", "));
      
        representationString += "]";
        return representationString;

    }

    private void toStringHelper(Node<K, V> root, ArrayList<Node<K, V>> nodes){
        if(root != null){
            toStringHelper(root.left, nodes);
            nodes.add(root);
            toStringHelper(root.right, nodes);
        }
    }

    public static void main(String[] args) {
        var bst = new BinarySearchTree<String, Integer>();
        String[] words = ("cash,cash,cash," + 
          "rules,rules,everything," + 
          "around,around,me," +
          "me,me,dollar," +
          "dollar,dollar,bill," + 
          "bill,okay,yes," + 
          "lucy,lucy,ricardo," + 
          "desi,desi,things")
          .split(",");
        for (String word : words) {
            Integer count = bst.lookup(word);
            bst.add(word, count == null ? 1 : count + 1);
        }

        System.out.print(bst);
    }
}