二叉树前中后序遍历递归转循环

Posted 2020-11-26 torapture

通过观察递归实现，用循环和栈模拟递归实现中结点入栈和出栈的过程。

#include <bits/stdc++.h>
#define DBG(x) cerr << #x << " = " << x << endl

using namespace std;
typedef long long LL;

struct Node {
    int val;
    Node *left, *right;

    Node() : left(NULL), right(NULL) {}

    Node(int val) : val(val), left(NULL), right(NULL) {}
};

enum StackState {
    AFTER_PUSH_ROOT = 0,
    AFTER_PUSH_LEFT = 1,
    AFTER_PUSH_RIGHT = 2,
};

void dfs_pre_order(Node *root, vector<int> &ret) {
    ret.push_back(root->val);
    if (root->left) dfs_pre_order(root->left, ret);
    if (root->right) dfs_pre_order(root->right, ret);
}

void dfs_in_order(Node *root, vector<int> &ret) {
    if (root->left) dfs_in_order(root->left, ret);
    ret.push_back(root->val);
    if (root->right) dfs_in_order(root->right, ret);
}

void dfs_post_order(Node *root, vector<int> &ret) {
    if (root->left) dfs_post_order(root->left, ret);
    if (root->right) dfs_post_order(root->right, ret);
    ret.push_back(root->val);
}

void pre_order(Node *root, vector<int> &ret) {
    stack<pair<Node *, StackState>> stk;
    stk.push(pair<Node *, StackState>(root, AFTER_PUSH_ROOT));
    while (!stk.empty()) {
        pair<Node *, StackState> &top = stk.top();
        switch (top.second) {
            case AFTER_PUSH_ROOT:
                ret.push_back(top.first->val);
                if (top.first->left) stk.push(pair<Node *, StackState>(top.first->left, AFTER_PUSH_ROOT));
                top.second = AFTER_PUSH_LEFT;
                break;
            case AFTER_PUSH_LEFT:
                if (top.first->right) stk.push(pair<Node *, StackState>(top.first->right, AFTER_PUSH_ROOT));
                top.second = AFTER_PUSH_RIGHT;
                break;
            case AFTER_PUSH_RIGHT:
                stk.pop();
                break;
        }
    }
}

void in_order(Node *root, vector<int> &ret) {
    stack<pair<Node *, StackState>> stk;
    stk.push(pair<Node *, StackState>(root, AFTER_PUSH_ROOT));
    while (!stk.empty()) {
        pair<Node *, StackState> &top = stk.top();
        switch (top.second) {
            case AFTER_PUSH_ROOT:
                if (top.first->left) stk.push(pair<Node *, StackState>(top.first->left, AFTER_PUSH_ROOT));
                top.second = AFTER_PUSH_LEFT;
                break;
            case AFTER_PUSH_LEFT:
                ret.push_back(top.first->val);
                if (top.first->right) stk.push(pair<Node *, StackState>(top.first->right, AFTER_PUSH_ROOT));
                top.second = AFTER_PUSH_RIGHT;
                break;
            case AFTER_PUSH_RIGHT:
                stk.pop();
                break;
        }
    }
}

void post_order(Node *root, vector<int> &ret) {
    stack<pair<Node *, StackState>> stk;
    stk.push(pair<Node *, StackState>(root, AFTER_PUSH_ROOT));
    while (!stk.empty()) {
        pair<Node *, StackState> &top = stk.top();
        switch (top.second) {
            case AFTER_PUSH_ROOT:
                if (top.first->left) stk.push(pair<Node *, StackState>(top.first->left, AFTER_PUSH_ROOT));
                top.second = AFTER_PUSH_LEFT;
                break;
            case AFTER_PUSH_LEFT:
                if (top.first->right) stk.push(pair<Node *, StackState>(top.first->right, AFTER_PUSH_ROOT));
                top.second = AFTER_PUSH_RIGHT;
                break;
            case AFTER_PUSH_RIGHT:
                ret.push_back(top.first->val);
                stk.pop();
                break;
        }
    }
}

Node *get_tree(int num) {
    Node *root = new Node(random());
    vector<Node *> nodes;
    nodes.push_back(root);

    for (int i = 0; i < num; i++) {
        unsigned int pos = random() * random() * random() % nodes.size();
        if (nodes[pos]->left == NULL && nodes[pos]->right == NULL) {
            if (random() % 2 == 0) {
                nodes[pos]->left = new Node(random());
                nodes.push_back(nodes[pos]->left);
            } else {
                nodes[pos]->right = new Node(random());
                nodes.push_back(nodes[pos]->right);
            }
        } else if (nodes[pos]->left) {
            nodes[pos]->right = new Node(random());
            nodes.push_back(nodes[pos]->right);
        } else {
            nodes[pos]->left = new Node(random());
            nodes.push_back(nodes[pos]->left);
        }
        if (nodes[pos]->left && nodes[pos]->right) {
            nodes.erase(nodes.begin() + pos);
        }
    }

    return root;
}

void free_tree(Node *root) {
    if (root->left) free_tree(root->left);
    if (root->right) free_tree(root->right);
    delete root;
}

void test_pre_order() {
    Node *tree = get_tree(100000);
    vector<int> dfs_ret, ret;

    dfs_pre_order(tree, dfs_ret);
    pre_order(tree, ret);
    assert(dfs_ret == ret);

    free_tree(tree);
}

void test_in_order() {
    Node *tree = get_tree(100000);
    vector<int> dfs_ret, ret;

    dfs_in_order(tree, dfs_ret);
    in_order(tree, ret);
    assert(dfs_ret == ret);

    free_tree(tree);
}

void test_post_order() {
    Node *tree = get_tree(100000);
    vector<int> dfs_ret, ret;

    dfs_post_order(tree, dfs_ret);
    post_order(tree, ret);
    assert(dfs_ret == ret);

    free_tree(tree);
}

int main(int argc, char **argv) {
    srand(time(NULL));

    test_pre_order();
    test_in_order();
    test_post_order();

    return 0;
}