C++: 使用Iterator遍历二叉树

上篇文章中我们实现了一个泛型的Find函数，并实现了支持这个泛型函数的链表Iterator, 这次，我们来实现支持这个函数的二叉树Iterator:

定义二叉树节点

template <typename T>
class tree_node
{
public:
    T value;
    tree_node* left;
    tree_node* right;

    explicit tree_node(T t=0)
    {
        value = t;
        left = nullptr;
        right = nullptr;
    }

    bool operator == (const T t) const
    {
        return value == t;
    }
    bool operator == (const tree_node& n) const
    {
        return this->value == n.value;
    }
    bool operator != (const T t) const
    {
        return value != t;
    }
    bool operator != (const tree_node& n) const
    {
        return this->value != n.value;
    }
};

同样，定义了一个模板类作为二叉树的节点

二叉树Iterator基类

由于二叉树有先序，中序，后序三种遍历方式，所以我们可以定义一个基类，声明好接口：

template <typename node>
class tree_iterator
{
protected:
    node* curr_ptr;
    node* root_ptr;
    std::stack<node*> st;

public:
    tree_iterator()
    {
        curr_ptr = nullptr;
        root_ptr = nullptr;
    }
    explicit tree_iterator(node* ptr)
    {
        curr_ptr = ptr;
        root_ptr = ptr;
    }

    virtual tree_iterator<node>& operator ++() = 0;

    tree_iterator<node>& operator ++(int)
    {
        return operator++();
    }

    node& operator * ()
    {
        return *curr_ptr;
    }

    node* operator -> ()
    {
        return curr_ptr;
    }

    node* get() const
    {
        return curr_ptr;
    }

    const node& operator * () const
    {
        return curr_ptr->value;
    }

    bool operator == (const tree_iterator& src)
    {
        return curr_ptr == src.get();
    }

    bool operator != (const tree_iterator& src)
    {
        return curr_ptr != src.get();
    }
};

上面的tree_iterator实现了基本的接口，但是把重载++声明为纯虚函数，下放给子类去实现

三种遍历方式

不同的遍历Iterator只需要实现不同的++即可

1.先序遍历

template <typename node>
class pre_order_tree_iterator : public tree_iterator<node>
{
public:
    using tree_iterator<node>::operator++;

    explicit pre_order_tree_iterator(node* p = nullptr) : tree_iterator<node>(p)
    {
        if(p)
            this->st.push(p);

        operator++();
    }

    tree_iterator<node> &operator++() override
    {
        auto& st = this->st;
        auto& curr_ptr = this->curr_ptr;

        if(st.empty())
            curr_ptr = nullptr;
        else
        {
            curr_ptr = st.top();
            st.pop();
            if (curr_ptr->right)
                st.push(curr_ptr->right);
            if(curr_ptr->left)
                st.push(curr_ptr->left);
        }
        return *this;
    }
};

2.中序遍历

template <typename node>
class in_order_tree_iterator : public tree_iterator<node>
{
public:
    using tree_iterator<node>::operator++;
    using tree_iterator<node>::st;
    using tree_iterator<node>::curr_ptr;

    explicit in_order_tree_iterator(node* p = nullptr) : tree_iterator<node>(p)
    {
        if(curr_ptr)
        {
            while (curr_ptr)
            {
                st.push(curr_ptr);
                curr_ptr = curr_ptr->left;
            }
            if(!st.empty())
            {
                curr_ptr = st.top();
                st.pop();
            }
        }
    }

    tree_iterator<node> &operator++() override
    {
        if(curr_ptr->right)
        {
            curr_ptr = curr_ptr->right;
            while(curr_ptr)
            {
                st.push(curr_ptr);
                curr_ptr = curr_ptr->left;
            }

        }
        if(!st.empty())
        {
            curr_ptr = st.top();
            st.pop();
        }
        else
            curr_ptr = nullptr;
        return *this;
    }
};

3.后序遍历

template <typename node>
class pos_order_tree_iterator : public tree_iterator<node>
{
public:
    using tree_iterator<node>::operator++;
    using tree_iterator<node>::st;
    using tree_iterator<node>::curr_ptr;

    explicit pos_order_tree_iterator(node* p = nullptr) : tree_iterator<node>(p)
    {
        if(curr_ptr)
        {
            while (curr_ptr)
            {
                st.push(curr_ptr);
                if(curr_ptr->right)
                    st.push(curr_ptr->right);
                curr_ptr = curr_ptr->left;
            }
            if(!st.empty())
            {
                curr_ptr = st.top();
                st.pop();
            }
        }
    }

    tree_iterator<node> &operator++() override
    {
        if(!st.empty())
        {
            if(st.top()->right == curr_ptr)/// 遍历过
            {
                curr_ptr = st.top();
                st.pop();
            }
            else/// 未遍历过
            {
                curr_ptr = st.top();
                st.pop();
                while(curr_ptr)
                {
                    st.push(curr_ptr);
                    if(curr_ptr->right)
                        st.push(curr_ptr->right);
                    curr_ptr = curr_ptr->left;
                }
                curr_ptr = st.top();
                st.pop();
            }
        }
        else
            curr_ptr = nullptr;
        return *this;
    }
};

测试

int main()
{
    auto n1 = new tree_node<int>(1);
    auto n2 = new tree_node<int>(2);
    auto n3 = new tree_node<int>(3);
    auto n4 = new tree_node<int>(4);
    auto n5 = new tree_node<int>(5);

    n1->left = n2;
    n1->right = n3;

    n2->left = n4;
    n2->right = n5;

    auto res = Find(pos_order_tree_iterator<tree_node<int>>(n1),
                     pos_order_tree_iterator<tree_node<int>>(),
                     1);

    if (res.get())
        cout << res->value << endl;
    else
        cout << "Not Found" << endl;

    return 0;
}