打印Leetcode二叉树

介绍

这个代码能够打印Leetcode中以字符串形式显示的二叉树，方便以直观的形式查看树的结构

代码

#include <iostream>
#include <vector>
#include <queue>
#include <cmath>
using namespace std;

/*
 * 此类用于将leetcode中string形式输出的二叉树实例化为树结构
 *    可进行如下string 形式的构造[1, 2, 3, 4, 5] 完全二叉树
 *                            [1,null,2,null,3,null,4,null]  单枝树                            
 */
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class binTree{
private:
    TreeNode *root;

    enum{NULLVAR = INT32_MIN}; 

    void freeNode(TreeNode *root){
        if(!root)
            return;
        freeNode(root->left);
        freeNode(root->right);
        delete root;
    };

    vector<int> extractIntArray(string str){
        vector<int> intList;
        str = str.substr(1, str.size() - 2);
        str += ',';
        for(int pos = 0; pos < str.size(); pos += 1){
            int tailPos = str.find(',', pos);
            string intStr = str.substr(pos, tailPos - pos);

            if(intStr == "null")
                intList.push_back(NULLVAR);
            else if(intStr == "")
                continue;
            else
                intList.push_back(stoi(intStr));
            pos = tailPos;
        }
        return intList;
    }

    void setBinTreeByInt(vector<int> &list){
        queue<TreeNode *> q;
        freeBinTree();

        if(list.empty() || list[0] == NULLVAR) // "[]" equal "[null]"
            return;

        // set root        
        root = new TreeNode(list[0]);
        q.push(root);
        for(int i = 1; i < list.size(); i++){
            TreeNode * parent = q.front();

            if(list[i] != NULLVAR){
                TreeNode *pnewnode = new TreeNode(list[i]);
                if(i % 2 != 0){
                    parent->left = pnewnode;
                }
                else{
                    parent->right = pnewnode;    
                }
                q.push(pnewnode);
            }
            if(i % 2 == 0)
                q.pop();
        }
    };

    vector<int> flattenBinTree(){
        vector<int> flatIntList;
        queue<TreeNode *> q;
        q.push(root);

        if(!root)
            return flatIntList;

        while(!q.empty()){
            TreeNode *ptn = q.front();
            
            if(ptn == NULL){
                flatIntList.push_back(NULLVAR);
            }
            else{
                flatIntList.push_back(ptn->val);
                q.push(ptn->left);
                q.push(ptn->right);
            }
            q.pop();
        }
        // erase back null
        while(flatIntList.back() == NULLVAR){
            flatIntList.pop_back();
        }
        return flatIntList;
    }

    vector<string> convertStrList(vector<int>& flatIntList){
        vector<string> flatStrList;
        for(auto var : flatIntList){
            if(var == NULLVAR)
                flatStrList.push_back("null");
            else
                flatStrList.push_back(to_string(var));
        }
        return flatStrList;
    }

    int treeHeight(TreeNode *root){
        if(root == NULL)
            return 0;
        return max(treeHeight(root->left), treeHeight(root->right)) + 1;
    }

    void writeRootToBoard(TreeNode *root,int height, int left, int right, vector<vector<string>> &board){
        if(root == NULL)
            return;
        int mid = (left + right) / 2;
        board[height][mid] = to_string(root->val);
        writeRootToBoard(root->left, height + 1, left, mid - 1, board);
        writeRootToBoard(root->right, height + 1, mid + 1, right, board);
    }

    vector<vector<string>> writeTreeToBoard(TreeNode* root, vector<vector<string>> &board){
        int treeheight = treeHeight(root);
        board = vector<vector<string>>(treeheight, vector<string>(pow(2, treeheight) - 1,""));

        writeRootToBoard(root, 0, 0, pow(2, treeheight) - 2, board);
        return board;
    }

public:
    

    binTree():root(NULL){};
    binTree(string str):root(NULL){
        setBinTree(str);
    };
    ~binTree(){
        freeNode(root);
    };
    void freeBinTree(){
        freeNode(root);
        root = NULL;
    }

    void setBinTree(string str){
        vector<int> intArray = extractIntArray(str);
        setBinTreeByInt(intArray);
    }

    string flatBinTreeToStr(){
        string flatStr = "[";
        vector<int> flatIntList = flattenBinTree();
        vector<string> flatStrList = convertStrList(flatIntList);
        if(flatStrList.empty())
            return "[]";
        
        for(auto str : flatStrList)
            flatStr += str + ',';
        flatStr[flatStr.size() - 1] = ']';
        
        return flatStr;
    }

    void printTreeToScreen(){
        vector<vector<string>> board;

        writeTreeToBoard(root, board);

        for(int i = 0; i < board.size(); i++){
            for(int j = 0; j < board[i].size(); j++){
                cout << '\t' << board[i][j];
            }
            cout << "\n\n";
        }
        return;
    }

    TreeNode * getTree(){
        return root;
    }
};

int main(){
    // 使用字符串初始化
    binTree bt("[1,2,3,null,4,null,null,5,6,null,7,8,null,null,null,null,null]");

    // 可以使用字符串初始化
    bt.setBinTree("[1,2,3]");

    // 返回树的根节点
    TreeNode * root = bt.getTree();

    // 反序列化为字符串
    string bt2str = bt.flatBinTreeToStr;
    cout << bt2str << endl;

    // 向屏幕打印字符串
    bt.printTreeToScreen();

}