TreeMap Structure and Visualization

The previous article Common Types of Binary Trees introduced binary search trees. Next, I will guide you to personally implement structures similar to Java's standard library TreeMap and TreeSet, aiding you in integrating knowledge with practice.

However, considering that this article is part of the data structure basics chapter, it will only explain the principles of TreeMap/TreeSet. The hands-on implementation of TreeMap/TreeSet is placed in Implementation of TreeMap/TreeSet after the Introduction to Binary Tree Exercises.

Unlike previous data structures like hash tables and queues, tree-related data structures require a strong understanding of recursion, making them more complex. If your grasp of recursion is not deep, the learning curve could be steep and not very meaningful. Even if you manage to understand it after great effort, you might still struggle with actual problems, which can be discouraging.

Therefore, I recommend a step-by-step approach. In the subsequent binary tree exercise chapter, over 100 algorithm problems will guide you in developing recursive thinking. After completing these exercises, you'll be able to tackle all binary tree-related algorithm problems easily, and understanding tree-related data structure implementations will seem simple. You may even find you don't need my code; you can implement TreeMap/TreeSet on your own intuition.

Alright, without further ado, let's get started.

Advantages of Binary Search Trees

In the previous article Common Types of Binary Trees, we introduced the characteristics of a Binary Search Tree (BST), namely "left smaller, right larger":

For each node in the tree, the value of every node in its left subtree must be smaller than the value of this node, and the value of every node in its right subtree must be larger than the value of this node.

For example, the tree below is a BST:

    7
   / \
  4   9
 / \   \
1   5   10

This characteristic of "left smaller, right larger" allows us to quickly find a specific node or all nodes within a certain range in a BST, which is the advantage of a BST.

You should have already learned about Binary Tree Traversal in the previous article. Below, we will use standard binary tree traversal functions combined with a visual panel to compare the operational differences between a BST and a regular binary tree.

You can expand the two panels below and click on the line if (targetNode !== null) to intuitively feel the efficiency difference between the two search algorithms:

The scenario displayed here is finding the target element. As you can see, by using the "left smaller, right larger" characteristic of a BST, you can quickly locate the target node. The ideal time complexity is the height of the tree, $O(logN)$, whereas a regular binary tree traversal function requires $O(N)$ time to traverse all nodes.

As for other operations like insertion, deletion, and modification, you first need to find the target node before performing these operations, correct? These operations merely involve changing a pointer, so their time complexity is also $O(logN)$.

Implementation Principles of TreeMap/TreeSet

By looking at the name TreeMap, you should be able to tell that its structure is similar to the previously introduced Hash Table HashMap. Both store key-value pairs. The HashMap stores key-value pairs in a table array, while TreeMap stores key-value pairs in the nodes of a binary search tree.

As for TreeSet, its relationship with TreeMap is similar to the relationship between HashMap and the hash set HashSet. Simply put, TreeSet is just a simple wrapper around TreeMap. Therefore, the following will mainly explain the implementation principles of TreeMap.

The classic TreeNode structure in LeetCode looks like this:

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { this.val = x; }
}

class TreeNode {
public:
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};

class TreeNode:
    def __init__(self, x: int):
        self.val = x
        self.left = None
        self.right = None

type TreeNode struct {
    Val   int
    Left  *TreeNode
    Right *TreeNode
}

var TreeNode = function(x) {
    this.val = x;
    this.left = null;
    this.right = null;
}

We can implement TreeMap by making a slight modification to the classic TreeNode structure:

// Capital K is the type for keys, capital V is the type for values
class TreeNode<K, V> {
    K key;
    V value;

    TreeNode<K, V> left;
    TreeNode<K, V> right;
    TreeNode(K key, V value) {
        this.key = key;
        this.value = value;
    }
}

// Capital K is the type for keys, capital V is the type for values
template <typename K, typename V>
class TreeNode {
public:
    K key;
    V value;
    TreeNode<K, V>* left;
    TreeNode<K, V>* right;
    TreeNode(K key, V value) : key(key), value(value), left(nullptr), right(nullptr) {}
};

class TreeNode:
    def __init__(self, key: K, value: V):
        self.key = key
        self.value = value
        self.left = None
        self.right = None

type TreeNode struct {
    Key   interface{}
    Value interface{}
    Left  *TreeNode
    Right *TreeNode
}

class TreeNode {
    constructor(key, value) {
        this.key = key;
        this.value = value;
        this.left = null;
        this.right = null;
    }
}

The TreeMap structure we are going to implement has the following API:

// Main interface of TreeMap
class MyTreeMap<K, V> {

    // ****** Basic methods for key-value mapping ******

    // Add/Update, complexity O(logN)
    public void put(K key, V value) {}

    // Get, complexity O(logN)
    public V get(K key) {}

    // Remove, complexity O(logN)
    public void remove(K key) {}

    // Check if contains key, complexity O(logN)
    public boolean containsKey(K key) {}

    // Return a collection of all keys in order, complexity O(N)
    public List<K> keys() {}

    // ****** Additional methods provided by TreeMap ******

    // Find the smallest key, complexity O(logN)
    public K firstKey() {}

    // Find the largest key, complexity O(logN)
    public K lastKey() {}

    // Find the largest key less than or equal to the given key, complexity O(logN)
    public K floorKey(K key) {}

    // Find the smallest key greater than or equal to the given key, complexity O(logN)
    public K ceilingKey(K key) {}

    // Find the key with rank k, complexity O(logN)
    public K selectKey(int k) {}

    // Find the rank of the key, complexity O(logN)
    public int rank(K key) {}

    // Range search, complexity O(logN + M), where M is the size of the range
    public List<K> rangeKeys(K low, K high) {}
}

In addition to the standard methods for adding, deleting, and querying such as get, put, remove, and containsKey, TreeMap offers many extra methods mainly related to the order of keys. Impressive, isn't it?

Hash tables are practical, but they struggle with handling the order of keys effectively. The LinkedHashMap implemented in the previous article Enhancing Hash Tables with Linked Lists only arranges keys by the order of insertion, and still cannot sort them by order of size.