支持存储自定义数据 的 多叉树 实现[c++]

#ifndef MTREE_H
#define MTREE_H

#include <algorithm>
#include <iostream>
#include <list>
#include <unordered_map>
#include <cassert>

/*!
 * \brief 定义 MTreeInfo 结构体
 * \param MTreeFlag 树节点的唯一标识类型
 * \param MTreeData 树节点的自定义数据类型
 */
template <typename MTreeFlag, typename MTreeData>
struct MTreeInfo {
    MTreeFlag selfid;
    MTreeFlag parentid;
    MTreeData* data;
};

/*!
 * \brief 定义树节点结构
 * \param MTreeFlag 树节点的唯一标识类型
 * \param MTreeData 树节点的自定义数据类型
 */
template <typename MTreeFlag, typename MTreeData>
struct MTreeNode {
    MTreeFlag id;
    MTreeFlag parent_id;
    std::list<MTreeNode*> children;
    MTreeNode* parent = nullptr;

    MTreeNode() = default;
    MTreeNode(MTreeFlag id, MTreeFlag parent_id)
        : id(id), parent_id(parent_id) {}
    ~MTreeNode() {
        for (MTreeNode* child : children) {
            delete child;
        }
    }
    // 深拷贝
    MTreeNode(const MTreeNode& other)
        : id(other.id), parent_id(other.parent_id) {
        for (const MTreeNode* child : other.children) {
            children.push_back(new MTreeNode(*child));
        }
    }

    // 浅拷贝
    MTreeNode& operator=(const MTreeNode& other) {
        if (this != &other) {
            id = other.id;
            parent_id = other.parent_id;
        }
        return *this;
    }

    void AddChild(MTreeNode* child) {
        children.push_back(child);
        child->parent = this;
    }

    void RemoveChild(MTreeNode* child) {
        auto it = std::find(children.begin(), children.end(), child);
        if (it != children.end()) {
            children.erase(it);
            child->parent = nullptr;
        }
    }
};

/*!
 * \brief 多叉树
 * \param MTreeFlag 树节点的唯一标识类型
 * \param MTreeData 树节点的自定义数据类型
 */
template<typename MTreeFlag, typename MTreeData>
class MTree {
private:
    std::unordered_map<MTreeFlag, MTreeNode<MTreeFlag, MTreeData>*> nodes_;
    MTreeNode<MTreeFlag, MTreeData> root_;

public:
    MTree() { root_.parent = nullptr; }
    ~MTree() = default;
    // 深拷贝
    MTree(const MTree& other) {
        for (const auto& pair : other.nodes_) {
            auto node = new MTreeNode<MTreeFlag, MTreeData>(*pair.second);
            nodes_[node->id] = node;
            node->parent = (pair.second->parent != nullptr)
                    ? nodes_[pair.second->parent->id]
                    : nullptr;
            for (MTreeNode<MTreeFlag, MTreeData>* chd : pair.second->children) {
                node->AddChild(new MTreeNode<MTreeFlag, MTreeData>(*chd));
            }
        }
    }

    // 浅拷贝
    MTree& operator=(const MTree& other) {
        if (this == &other)
            return *this;

        CleanUp();
        for (const auto& pair : other.nodes_) {
            auto node = new MTreeNode<MTreeFlag, MTreeData>(*pair.second);
            nodes_[node->id] = node;
            node->parent = (pair.second->parent != nullptr)
                    ? nodes_[pair.second->parent->id]
                    : nullptr;
            for (MTreeNode<MTreeFlag, MTreeData>* chd : pair.second->children) {
                node->AddChild(new MTreeNode<MTreeFlag, MTreeData>(*chd));
            }
        }
        return *this;
    }

    // 创建树
    void CreateTree(const std::list<MTreeInfo<MTreeFlag, MTreeData>>& info_list) {
        for (const auto& info : info_list) {
            if (nodes_.find(info.selfid) == nodes_.end()) {
                auto node = new MTreeNode<MTreeFlag, MTreeData>(info.selfid, info.parentid);
                if (info.parentid == 0) {
                    // 如果是根节点
                    root_.AddChild(node);
                } else {
                    // 否则，找到父节点并将其作为子节点
                    auto parent_node_iter = nodes_.find(info.parentid);
                    if (parent_node_iter != nodes_.end()) {
                        parent_node_iter->second->AddChild(node);
                    }
                }
                nodes_[node->id] = node; // 添加到节点映射中
            }
        }
    }

    // 获取根节点
    MTreeNode<MTreeFlag, MTreeData>* GetRoot() const {
        return &root_;
    }

    // 根据 ID 获取节点的指针
    MTreeNode<MTreeFlag, MTreeData>* GetNodeById(MTreeFlag id) const {
        auto it = nodes_.find(id);
        return (it != nodes_.end()) ? it->second : nullptr;
    }

    // 移除指定 ID 的节点，并更新树结构
    void RemoveNodeById(MTreeFlag id) {
        auto it = nodes_.find(id);
        if (it != nodes_.end()) {
            MTreeNode<MTreeFlag, MTreeData>* node_to_remove = it->second;
            if (node_to_remove->parent) {
                node_to_remove->parent->RemoveChild(node_to_remove);
            }
            delete node_to_remove; // 删除节点
            nodes_.erase(it); // 从映射中移除节点
        }
    }

    // 获取节点的兄弟节点
    std::list<MTreeNode<MTreeFlag, MTreeData>*>
    GetSiblings(MTreeNode<MTreeFlag, MTreeData>* node) const {
        if (!node || node->parent == nullptr) {
            return {};
        }
        std::list<MTreeNode<MTreeFlag, MTreeData>*> siblings;
        for (MTreeNode<MTreeFlag, MTreeData>* sibling : node->parent->children) {
            if (sibling != node) {
                siblings.push_back(sibling);
            }
        }
        return siblings;
    }

    // 清理树内存
    void CleanUp() {
        for (auto& pair : nodes_) {
            delete pair.second;
        }
        nodes_.clear();
    }

    // 根据 MTreeInfo 结构体更新或新增节点
    void UpdateOrAddNode(const MTreeInfo<MTreeFlag, MTreeData>& info) {
        auto it = nodes_.find(info.selfid);
        if (it != nodes_.end()) {
            // 如果节点已存在，则更新父节点
            MTreeNode<MTreeFlag, MTreeData>* node = it->second;
            if (node->parent_id != info.parentid) {
                node->parent->RemoveChild(node);
                node->parent_id = info.parentid;

                // 如果新父节点 ID 为 0，则将节点挂载到根节点下
                if (info.parentid == 0) {
                    root_.AddChild(node);
                } else {
                    // 否则，找到父节点并将其作为子节点
                    auto parent_node_iter = nodes_.find(info.parentid);
                    if (parent_node_iter != nodes_.end()) {
                        parent_node_iter->second->AddChild(node);
                    }
                }
            }
        } else {
            // 如果节点不存在，则创建新节点并添加到树中
            auto node = new MTreeNode<MTreeFlag, MTreeData>(info.selfid, info.parentid);
            if (info.parentid == 0) {
                // 如果新父节点 ID 为 0，则将节点挂载到根节点下
                root_.AddChild(node);
            } else {
                // 否则，找到父节点并将其作为子节点
                auto parent_node_iter = nodes_.find(info.parentid);
                if (parent_node_iter != nodes_.end()) {
                    parent_node_iter->second->AddChild(node);
                }
            }
            nodes_[node->id] = node; // 更新节点映射
        }
    }

    // 判断已知节点是否存在于树中的接口, 不包括根节点
    bool ContainsNode(MTreeNode<MTreeFlag, MTreeData>* node) const {
        assert(node);

        // 根节点不包含在内
        if (node == &this->root) {
            return false;
        }

        // 父节点不存在
        if (node->parent && nodes_.find(node->parent->id) == nodes_.end()) {
            return false;
        }

        // 如果节点有父节点，检查节点是否在父节点的子节点列表中
        auto& siblings = node->parent->children;
        auto it = std::find(siblings.begin(), siblings.end(), node);
        if (it == siblings.end()) {
            return false; // 节点不在父节点的子节点列表中
        }

        return true;
    }

    // 索引 已知node在tree中的位置, 并为其添加subnode
    void AddSubNode(MTreeNode<MTreeFlag, MTreeData>* parent,
                 MTreeNode<MTreeFlag, MTreeData>* child) {
        assert(parent);
        assert(child);

        if (!ContainsNode(parent))
            return;

        parent->AddChild(child);
        child->parent = parent;
        nodes_[child->id] = child;
    }
};

#endif // MTREE_H