使用 Java 手动实现 TreeMap

2024-01-17 04:13:34

前言

TreeMap 是 Java 中的一个数据结构，它允许您存储键值对，并根据键对值进行排序。这使得它非常适合存储和检索数据，例如用户数据或产品目录。

基本原理

TreeMap 使用红黑树来实现。红黑树是一种自平衡二叉搜索树，它可以保证在 O(log n) 的时间复杂度内进行查找、插入和删除操作。

手写实现

public class TreeMap<K extends Comparable<K>, V> {

    private Node<K, V> root;

    public void put(K key, V value) {
        root = put(root, key, value);
    }

    private Node<K, V> put(Node<K, V> node, K key, V value) {
        if (node == null) {
            return new Node<>(key, value);
        }

        int cmp = key.compareTo(node.key);
        if (cmp == 0) {
            node.value = value;
        } else if (cmp < 0) {
            node.left = put(node.left, key, value);
        } else {
            node.right = put(node.right, key, value);
        }

        return node;
    }

    public V get(K key) {
        Node<K, V> node = get(root, key);
        return node == null ? null : node.value;
    }

    private Node<K, V> get(Node<K, V> node, K key) {
        if (node == null) {
            return null;
        }

        int cmp = key.compareTo(node.key);
        if (cmp == 0) {
            return node;
        } else if (cmp < 0) {
            return get(node.left, key);
        } else {
            return get(node.right, key);
        }
    }

    public void remove(K key) {
        root = remove(root, key);
    }

    private Node<K, V> remove(Node<K, V> node, K key) {
        if (node == null) {
            return null;
        }

        int cmp = key.compareTo(node.key);
        if (cmp == 0) {
            if (node.left == null) {
                return node.right;
            } else if (node.right == null) {
                return node.left;
            } else {
                Node<K, V> successor = findMin(node.right);
                node.key = successor.key;
                node.value = successor.value;
                node.right = remove(node.right, successor.key);
            }
        } else if (cmp < 0) {
            node.left = remove(node.left, key);
        } else {
            node.right = remove(node.right, key);
        }

        return node;
    }

    public K firstKey() {
        return firstKey(root);
    }

    private K firstKey(Node<K, V> node) {
        if (node == null) {
            return null;
        }

        if (node.left == null) {
            return node.key;
        } else {
            return firstKey(node.left);
        }
    }

    public K lastKey() {
        return lastKey(root);
    }

    private K lastKey(Node<K, V> node) {
        if (node == null) {
            return null;
        }

        if (node.right == null) {
            return node.key;
        } else {
            return lastKey(node.right);
        }
    }

    public int size() {
        return size(root);
    }

    private int size(Node<K, V> node) {
        if (node == null) {
            return 0;
        }

        return 1 + size(node.left) + size(node.right);
    }

    private static class Node<K, V> {

        private K key;
        private V value;
        private Node<K, V> left;
        private Node<K, V> right;

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

示例

以下是一个示例，展示如何使用 TreeMap 来存储和检索数据：

import java.util.TreeMap;

public class Main {

    public static void main(String[] args) {
        // 创建一个 TreeMap
        TreeMap<String, String> map = new TreeMap<>();

        // 向 TreeMap 中添加一些数据
        map.put("John", "123-456-7890");
        map.put("Mary", "987-654-3210");
        map.put("Bob", "456-789-0123");

        // 从 TreeMap 中获取数据
        String phoneNumber = map.get("John");
        System.out.println(phoneNumber); // 输出：123-456-7890

        // 遍历 TreeMap
        for (String key : map.keySet()) {
            String value = map.get(key);
            System.out.println(key + ": " + value);
        }
    }
}