二叉树展开为链表

• 迭代实现 先序遍历，同时进行链表链接操作。

题目描述

给你二叉树的根结点 root ，请你将它展开为一个单链表：

展开后的单链表应该同样使用 TreeNode ，其中 right 子指针指向链表中下一个结点，而左子指针始终为 null 。 展开后的单链表应该与二叉树 先序遍历 顺序相同。

示例：

输入：root = [1,2,5,3,4,null,6]
输出：[1,null,2,null,3,null,4,null,5,null,6]

C++实现

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    void flatten(TreeNode* root) {
        // 1. 迭代先序遍历 -> vector
        // 2. vector 转为链表
        if (!root) {
            return;
        }

        vector<TreeNode *> vec;
        stack<TreeNode *> stk;
        stk.push(root);
        TreeNode *prev = nullptr;
        while (!stk.empty()) {
            auto q = stk.top();
            stk.pop();

            if (prev) {
                prev->right = q;
                prev->left = nullptr;
            }

            if (q->right) {
                stk.push(q->right);
            }
            if (q->left) {
                stk.push(q->left);
            }
            prev = q;
        }

        // int n = vec.size();
        // for (int i = 1; i < n; i++) {
        //     vec[i-1]->right = vec[i];
        //     vec[i-1]->left = nullptr;
        // }
        // vec[n-1]->right = nullptr;
        // vec[n-1]->left = nullptr;

        // TreeNode *prev = nullptr;
        // for (auto & e : vec) {
        //     if (prev) {
        //         prev->right = e;
        //         prev->left = nullptr;
        //     }
        //     prev = e;
        // }
        // prev->left = nullptr;
        // prev->right = nullptr;
    }
};

golang实现

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func flatten(root *TreeNode)  {
    // 迭代 + 转为链表
    if root == nil {
        return
    }

    var prev *TreeNode
    stk := []*TreeNode{root}
    for len(stk) > 0 {
        n := len(stk)
        curr := stk[n-1]
        stk = stk[:n-1]
        if prev != nil {
            prev.Right = curr
            prev.Left = nil
        }

        if curr.Right != nil {
            stk = append(stk, curr.Right)
        }
        if curr.Left != nil {
            stk = append(stk, curr.Left)
        }
        prev = curr
    }
}