TreeMap Structure and Visualization

The previous article Common Types of Binary Trees introduces binary search trees. Next, I will guide you through implementing a TreeMap and TreeSet structure similar to those in the Java standard library, to help you combine theory with practice.

However, considering this article is still within the foundational data structures chapter, we will only explain the principles of TreeMap/TreeSet. The hands-on implementation of TreeMap/TreeSet is placed at the end of the binary tree exercise series: 动手实现 TreeMap/TreeSet.

Unlike earlier data structures like hash tables and queues, tree-related structures require a strong understanding of recursion, pushing the difficulty to a higher level. If your grasp of recursion is not deep enough, discussing it now might not only steepen the learning curve but also be of little benefit. Even if you manage to understand it after much effort, you might still struggle with practical problems, which could be discouraging.

Therefore, I recommend a gradual approach. In the upcoming binary tree exercises chapter, we will guide you through over 100 algorithm problems to develop your recursive thinking. Once you complete these, you'll be able to tackle all binary tree-related algorithm problems effortlessly. Implementing tree-related data structures will then feel very straightforward. You might not even need to refer to my code and 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.