[C++] 深入理解红黑树:封装set与map

1. 源码和框架

GCC 15.2.0版本中,map和set的源代码在map/set/stl_map.h/stl_set.h/stl_tree.h等⼏个头⽂件里。通过阅读源代码我们可以大致理解,set和map的底层套用了同一个红黑树 的模板。其中,为了兼容mapkey-value键值对,set内部将key同时命名为keyvalue,以便于后续调用红黑树统一操作,实现了代码复用。以下是源码的大致框架:

网址 Set&MapOrigCodeFrameWork

2. 模拟实现set和map

2.1 结构和功能

这里主要实现的是基于同一个红黑树模板的set和map,包含以下功能:

  1. 正向iterator和const_iterator
  2. 插入和查找操作

代码结构如下:

2.2 代码细节介绍

部分代码在源代码的基础上进行了微调:

  1. key的参数设置为K,value的参数设置为V,红黑树使用T(type)作为数据类型。
  2. 由于红黑树需要同时支持set和map的操作,因此在insert内部我们无法直接比较,需要实现:容器为set时比较key ,容器为map时取出pair中的key 进行比较。因此我们在set和map层分别实现一个仿函数,用于取出key进行比较。
  3. set的迭代器不支持修改key,需要把第二个模板参数改成const K;而map的迭代器不支持修改key但支持修改value,因此第二个参数需设置为pair<const K, V>

除此之外,代码还支持iterator的实现,其大致框架和list的iterator是一致的。即用类型封装节点的指针,再通过重载运算符达到迭代器像指针一样访问的行为。重点在于operator++operator--的实现:

2.1 重载operator++

之前我们在学习如何使用set的过程中,我们得知迭代器走的是中序遍历(左->根->右) ,因此在重载++时,核心逻辑就是只看局部,只考虑当前中序需要访问的下一个节点。分类讨论如下:

  1. 当右子树存在时,下一个节点为 子树的最节点;
  2. 当右子树不存在时,向上遍历直至该节点为左孩子,下一个节点即为该节点的父节点。特殊情况 :如果找不到符合标准的节点(遍历到根,此时parent为nullptr,cur为root),说明下一个节点为end(),此时赋值parent空指针也是符合结果的,因此归为一类。

代码如下:

cpp 复制代码
//重载++
Self& operator++()
{
    //1.右不为空,则下一个节点为右子树的最左节点
    if (_node->_right)
    {
        Node* leftmost = _node->_right;
        while (leftmost && leftmost->_left)
        {
            leftmost = leftmost->_left;
        }

        _node = leftmost;
    }
    //2.右为空,向上查找直至该节点为左孩子时,返回父节点
    else
    {
        Node* cur = _node;
        Node* parent = _node->_parent;
        while (parent && cur == parent->_right)
        {
            cur = parent;
            parent = parent->_parent;
        }

        _node = parent; //包含了最后一个节点的特殊情况
    }

    return *this;
}

2.2 重载operator--

operator--operator++的逻辑类似,反过来即可。代码实现如下:

cpp 复制代码
//重载--
Self& operator--()
{
    //1.节点为end(),则找树的最右侧节点
    if (_node == nullptr)
    {
        Node* rightmost = _root;
        while (rightmost && rightmost->_right)
        {
            rightmost = rightmost->_right;
        }

        _node = rightmost;
    }
    //2.左子树不为空,找左子树的最右节点
    else if (_node->_left)
    {
        Node* rightmost = _node->_left;
        while (rightmost->_right)
        {
            rightmost = rightmost->_right;
        }

        _node = rightmost;
    }
    //2.左子树为空,向上查找直至该节点为右孩子时,返回父节点
    else
    {
        Node* cur = _node;
        Node* parent = _node->_parent;
        while (parent && cur != parent->_right)
        {
            cur = parent;
            parent = parent->_parent;
        }

        _node = parent;
    }

    return *this;
}

2.3 map中operator 的实现

map是支持通过key查找value的,而实现这个功能还需要insert返回值的支持:pair<Iterator, bool> Insert(const T& data)

通过pair中的iterator即可获取对应的value。

cpp 复制代码
//已知key存在,则可以通过插入的返回值反向找到value
V& operator[](const K& key)
{
    pair<iterator, bool> ret = _t.Insert({ key, V() }); //V()表示value的缺省值
	  return (ret.first)->second; //返回值是iterator的value
}
cpp 复制代码
pair<iterator, bool> insert(const pair<K, V>& kv)
{
		return _t.Insert(kv);
}

3. 代码实现

3.1 mySet.h

cpp 复制代码
#pragma once
#include<utility>
#include "RBTree.h"
using namespace std;


namespace mzh
{
	template<class K>
	class set
	{
		//仿函数SetKeyOfT,用于取出set中的key或者map中的pair.first
		struct SetKeyOfT
		{
			const K& operator()(const K& key) const
			{
				return key;
			}
		};

	public:
		//红黑树迭代器中的类型名,加typename保证语法正确,设置const K以免key被修改
		typedef typename RBTree<K, const K, SetKeyOfT>::Iterator iterator;
		typedef typename RBTree<K, const K, SetKeyOfT>::ConstIterator const_iterator;

		//迭代器begin
		iterator begin()
		{
			return _t.Begin();
		}

		//迭代器end
		iterator end()
		{
			return _t.End();
		}

		//const迭代器begin
		const_iterator begin() const
		{
			return _t.Begin();
		}

		//const迭代器end
		const_iterator end() const
		{
			return _t.End();
		}

		//插入
		pair<iterator, bool> insert(const K& key)
		{
			//_t.Insert(key);
			return _t.Insert(key);
		}

		//查找
		iterator find(const K& key)
		{
			return _t.Find(key);
		}
		 
	private:
		RBTree<K, const K, SetKeyOfT> _t; //加const避免插入的返回值被修改
	};
}

3.2 myMap.h

cpp 复制代码
#pragma once
#include<utility>
#include "RBTree.h"
using namespace std;

namespace mzh
{
	template<class K, class V>
	class map
	{
		struct MapKeyOfT
		{
			const K& operator()(const pair<K, V>& kv)
			{
				return kv.first;
			}
		};

	public:
		typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Iterator iterator;
		typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::ConstIterator const_iterator;

		iterator begin()
		{
			return _t.Begin();
		}

		iterator end()
		{
			return _t.End();
		}

		const_iterator begin() const
		{
			return _t.Begin();
		}

		const_iterator end() const
		{
			return _t.End();
		}

		pair<iterator, bool> insert(const pair<K, V>& kv)
		{
			return _t.Insert(kv);
		}

		iterator find(const K& key)
		{
			return _t.Find(key);
		}

		//出函数后,返回对象未被销毁时才可以返回引用
		//已知key存在,则可以通过插入的返回值反向找到value
		V& operator[](const K& key)
		{
			pair<iterator, bool> ret = _t.Insert({ key, V() }); //V()表示value的缺省值
			return (ret.first)->second; //返回值是iterator的value
		}

	private:
		//底层红黑树:键的类型,节点存储的数据类型,键萃取器(从pair中取出key)
		RBTree<K, pair<const K, V>, MapKeyOfT> _t;
	};
}

3.3 RBTree.h

cpp 复制代码
#pragma once
#include<utility>
#include<iostream>
using namespace std;

namespace mzh
{
    //枚举值表示颜色
    enum Colour
    {
        RED,
        BLACK
    };

    //节点
    template<class T>
    struct RBTreeNode
    {
    public:
        T _data; //让节点可以同时用于set(key)和map(pair)
        RBTreeNode<T>* _left;
        RBTreeNode<T>* _right;
        RBTreeNode<T>* _parent;
        Colour _col;

        //构造函数
        RBTreeNode(const T& data)
            :_data(data)
            , _left(nullptr)
            , _right(nullptr)
            , _parent(nullptr)
            , _col(RED)
        {
        }
    };

    //迭代器
    template<class T, class Ref, class Ptr>
    struct RBTreeIterator
    {
    public:
        typedef RBTreeNode<T> Node; //节点类型
        typedef RBTreeIterator<T, T&, T*> Self; //红黑树迭代器(根据RBTreeIterator的模板)

        Node* _node;
        Node* _root;

        //构造
        RBTreeIterator(Node* node, Node* root)
            :_node(node)
            , _root(root)
        {
        }

        //重载++
        Self& operator++()
        {
            //1.右不为空,则下一个节点为右子树的最左节点
            if (_node->_right)
            {
                Node* leftmost = _node->_right;
                while (leftmost && leftmost->_left)
                {
                    leftmost = leftmost->_left;
                }

                _node = leftmost;
            }
            //2.右为空,向上查找直至该节点为左孩子时,返回父节点
            else
            {
                Node* cur = _node;
                Node* parent = _node->_parent;
                while (parent && cur == parent->_right)
                {
                    cur = parent;
                    parent = parent->_parent;
                }

                _node = parent;
            }

            return *this;
        }

        //重载--
        Self& operator--()
        {
            //1.节点为end(),则找树的最右侧节点
            if (_node == nullptr)
            {
                Node* rightmost = _root;
                while (rightmost && rightmost->_right)
                {
                    rightmost = rightmost->_right;
                }

                _node = rightmost;
            }
            //2.左子树不为空,找左子树的最右节点
            else if (_node->_left)
            {
                Node* rightmost = _node->_left;
                while (rightmost->_right)
                {
                    rightmost = rightmost->_right;
                }

                _node = rightmost;
            }
            //2.左子树为空,向上查找直至该节点为右孩子时,返回父节点
            else
            {
                Node* cur = _node;
                Node* parent = _node->_parent;
                while (parent && cur != parent->_right)
                {
                    cur = parent;
                    parent = parent->_parent;
                }

                _node = parent;
            }

            return *this;
        }

        //重载*:返回整个对象(解引用)
        Ref operator*()
        {
            return _node->_data;
        }

        //重载->:返回对象内部成员的地址
        Ptr operator->()
        {
            return &_node->_data;
        }

        //重载!=
        bool operator!=(const Self& s) const
        {
            return _node != s._node;
        }

        //重载==
        bool operator==(const Self& s) const
        {
            return _node == s._node;
        }
    };

    //树
    template<class K, class T, class KeyOfT> //KeyOfT是仿函数类型
    class RBTree
    {
    private:
        typedef RBTreeNode<T> Node;
        Node* _root = nullptr;

    public:
        typedef RBTreeIterator<T, T&, T*> Iterator;
        typedef RBTreeIterator<T, const T&, const T*> ConstIterator;

        //迭代器Begin:找树的最左节点
        Iterator Begin()
        {
            Node* leftmost = _root;
            while (leftmost && leftmost->_left)
            {
                leftmost = leftmost->_left;
            }

            Iterator ret = Iterator(leftmost, _root);
            return ret; //或者直接返回构造return RBTreeIterator(leftmost, _root);
        }

        //迭代器End:根节点的_parent,即空指针
        Iterator End()
        {
            return Iterator(nullptr, _root);
        }

        //迭代器ConstBegin
        ConstIterator Begin() const
        {
            Node* leftmost = _root;
            while (leftmost && leftmost->_left)
            {
                leftmost = leftmost->_left;
            }

            return ConstIterator(leftmost, _root);
        }

        //迭代器ConstEnd
        ConstIterator End() const
        {
            return ConstIterator(nullptr, _root);
        }

        //析构函数
        ~RBTree()
        {
            Destroy(_root);
            _root = nullptr;
        }

        //插入
        pair<Iterator, bool> Insert(const T& data)
        {
            //根
            if (!_root)
            {
                _root = new Node(data);
                _root->_col = BLACK;
                return { Iterator(_root, _root), true }; //改成iterator的形式
            }

            //非根
            KeyOfT kot; //比较逻辑(可以像类一样使用)
            Node* cur = _root;
            Node* parent = nullptr;
            while (cur)
            {
                if (kot(cur->_data) < kot(data))
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else if (kot(cur->_data) > kot(data))
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else
                {
                    return { Iterator(cur, _root), false }; //改成iterator的形式
                }
            }

            //新建节点
            cur = new Node(data);
            cur->_col = RED;

            if (kot(parent->_data) < kot(data)) //右
            {
                parent->_right = cur;
            }
            else //左
            {
                parent->_left = cur;
            }

            cur->_parent = parent;

            //旋转调整
            while (parent && parent->_col == RED)
            {
                Node* grandfather = parent->_parent;

                //p在g左
                if (parent == grandfather->_left)
                {
                    Node* uncle = grandfather->_right;

                    //u存在且为红->变色继续向上处理
                    if (uncle && uncle->_col == RED)
                    {
                        parent->_col = BLACK;
                        uncle->_col = BLACK;
                        grandfather->_col = RED;

                        cur = grandfather;
                        parent = grandfather->_parent;
                    }
                    else //u不存在或为黑
                    {
                        if (cur == parent->_left)
                        {
                            RotateR(grandfather);

                            parent->_col = BLACK;
                            grandfather->_col = RED;
                        }
                        else
                        {
                            RotateL(parent);
                            RotateR(grandfather);

                            cur->_col = BLACK;
                            grandfather->_col = RED;
                        }

                        break;
                    }
                }
                //p在g右
                else
                {
                    Node* uncle = grandfather->_left;
                    //仅变色
                    if (uncle && uncle->_col == RED)
                    {
                        parent->_col = BLACK;
                        uncle->_col = BLACK;
                        grandfather->_col = RED;

                        cur = parent;
                        parent = grandfather;
                    }
                    else //旋转 + 变色
                    {
                        if (cur == parent->_right)
                        {
                            RotateL(grandfather);

                            parent->_col = BLACK;
                            grandfather->_col = RED;
                        }
                        else
                        {
                            RotateR(parent);
                            RotateL(grandfather);

                            cur->_col = BLACK;
                            grandfather->_col = RED;
                        }

                        break;
                    }
                }
            }

            //一律设置为黑
            _root->_col = BLACK;

            return { Iterator(cur, _root), true }; //改成iterator的形式
        }

        //通过key查找
        Iterator Find(const K& key)
        {
            KeyOfT kot; //比较逻辑(可以像类一样使用)
            Node* cur = _root;
            while (cur)
            {
                if (kot(cur->_data) > key)
                {
                    cur = cur->_left;
                }
                else if (kot(cur->_data) < key)
                {
                    cur = cur->_right;
                }
                else
                {
                    return Iterator(cur, _root);
                }
            }

            return End();
        }

    private:
        //右旋
        void RotateR(Node* pParent)
        {
            Node* parent = pParent;
            Node* subL = parent->_left;
            Node* subLR = subL->_right;

            //1.parent和subLR
            parent->_left = subLR;

            if (subLR)
            {
                subLR->_parent = parent;
            }

            Node* parentParent = parent->_parent;

            //2.parent和subL
            parent->_parent = subL;
            subL->_right = parent;

            //3.subL和parentParent
            if (parentParent == nullptr)
            {
                _root = subL;
                subL->_parent = nullptr;
            }
            else
            {
                if (parentParent->_left == parent)
                {
                    parentParent->_left = subL;
                }
                else
                {
                    parentParent->_right = subL;
                }

                subL->_parent = parentParent;
            }
        }

        //左旋
        void RotateL(Node* pParent)
        {
            Node* parent = pParent;
            Node* subR = parent->_right;
            Node* subRL = subR->_left;

            //1.parent和subRL
            parent->_right = subRL;
            if (subRL)
            {
                subRL->_parent = parent;
            }

            //2.parent和subR
            Node* parentParent = parent->_parent;

            parent->_parent = subR;
            subR->_left = parent;

            //3.subR和parentParent
            if (parentParent == nullptr)
            {
                _root = subR;
                subR->_parent = nullptr;
            }
            else
            {
                if (parentParent->_left == parent)
                {
                    parentParent->_left = subR;
                }
                else
                {
                    parentParent->_right = subR;
                }

                subR->_parent = parentParent;
            }
        }

        //递归销毁
        void Destroy(Node* root)
        {
            if (root == nullptr) return;

            Destroy(root->_left);
            Destroy(root->_right);
            delete root;
        }
    };
}

3.4 test_setmap.cpp

cpp 复制代码
#include<iostream>
#include<utility>
#include "myMap.h"
#include "mySet.h"
#include "RBTree.h"

using namespace std;
using namespace mzh;

void Print(const set<int>& s)
{
    set<int>::const_iterator it = s.end();
    while (it != s.begin())
    {
        --it;
        // 不⽀持修改
        //*it += 2;

        cout << *it << " ";
    }
    cout << endl;
}

void test_set()
{
    set<int> s;
    int a[] = { 4, 2, 6, 1, 3, 5, 15, 7, 16, 14 };
    int b[] = { 72, 19, 45, 8, 93, 31, 57, 62, 14, 87 };
    for (auto e : a)
    {
        s.insert(e);
    }
    for (auto e : b)
    {
        s.insert(e);
    }

    for (auto e : s)
    {
        cout << e << " ";
    }
    cout << endl;

    Print(s);
}

void test_map()
{
    map<string, string> dict;
    dict.insert({ "sort", "排序" });
    dict.insert({ "left", "左边" });
    dict.insert({ "right", "右边" });

    dict["left"] = "左边,剩余";
    dict["insert"] = "插⼊";
    dict["string"];

    map<string, string>::iterator it = dict.begin();
    while (it != dict.end())
    {
        // 不能修改first,可以修改second
        //it->first += 'x';
        it->second += " hehe";

        cout << it->first << ":" << it->second << endl;
        ++it;
    }
    cout << endl;
}


int main()
{
    cout << "test_set:" << endl;
    test_set();
    cout << endl;
    cout << "test_map:" << endl;
    test_map();
    cout << endl;
    cout << "end" << endl;
}
相关推荐
泡海椒1 小时前
JQuick-Curl 性能分析:并发场景下的性能表现与调优,第三方接口调用不只要快写也要稳跑
java·开发语言·okhttp
小智老师PMP1 小时前
2026深度解析|PMP第八版与NPDP核心侧重点本质区别(管理类证书怎么选)
开发语言·分布式·算法·职场和发展·产品经理
HugoStudio_SWAN1 小时前
洛谷 B4450 / B3867 / B3923 智慧购物、储蓄与做题——从程序到生活
c++·学习·程序人生·算法·生活
江畔柳前堤2 小时前
前台·中台·后台:2026年AI原生时代的架构全景图
开发语言·人工智能·算法·机器学习·架构·scala·ai-native
Figo_Cheung3 小时前
Figo基于RNC理论的宇宙演化第七纪元猜想
开发语言·php
会周易的程序员3 小时前
5Draft(五帝)测试报告
服务器·c++·分布式·raft·共识算法·共识·算力服务器
UIU1143 小时前
scanf与cout的误区:探究其内部的机制
c++·学习·c#·scanf
闻缺陷则喜何志丹3 小时前
【栈 运算系统】P3719 [AHOI2017初中组] rexp|普及+
c++·算法··洛谷·运算系统
泡海椒3 小时前
JQuick-Curl 二次开发:自定义解析器、扩展点开发指南,如何把框架能力真正变成团队资产
java·开发语言·okhttp·maven