二叉树展开为链表

题面

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

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

示例 1：

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

示例 2：

输入： root = [] 输出：[]

示例 3：

输入： root = [0] 输出：[0]

提示：

• 树中结点数在范围 [0, 2000] 内
• -100 <= Node.val <= 100

进阶： 你可以使用原地算法（O(1) 额外空间）展开这棵树吗？

分析

• 当一个节点已经遍历过时，其左右子树也都已经被加入到栈中，此时遍历到下一个节点时，可以直接将该节点赋值为上一个遍历的节点的右节点

• 代码实现

  /**
   * 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) {
          if (root == nullptr)
              return;
          TreeNode* dummy = new TreeNode(0);
          stack<TreeNode*> st;
          st.push(root);
          TreeNode* pre = dummy;
          while (!st.empty()) {
              TreeNode* top = st.top();
              st.pop();
              pre->left = nullptr;
              pre->right = top;
              pre = top;
              if (top->right)
                  st.push(top->right);
              if (top->left)
                  st.push(top->left);
          }
          delete dummy;
      }
  };

• 显然，这个方法的空间复杂度是 $O(n)$，因为前序遍历过程中需要使用栈存储节点，那么如何只在 $O(1)$ 的空间复杂度下实现？

• 注意到前序遍历访问各节点的顺序是根节点、左子树、右子树。如果一个节点的左子节点为空，则该节点不需要进行展开操作

• 如果一个节点的左子节点不为空，则该节点的左子树中的最后一个节点被访问之后，该节点的右子节点被访问

• 该节点的左子树中最后一个被访问的节点是左子树中的最右边的节点，也是该节点的前驱节点。因此，问题转化成寻找当前节点的前驱节点

• 具体做法是，对于当前节点，如果其左子节点不为空，则在其左子树中找到最右边的节点，作为前驱节点，将当前节点的右子节点赋给前驱节点的右子节点，然后将当前节点的左子节点赋给当前节点的右子节点，并将当前节点的左子节点设为空

• 对当前节点处理结束后，继续处理链表中的下一个节点，直到所有节点都处理结束

• 代码实现如下

  class Solution {
  public:
      void flatten(TreeNode* root) {
          TreeNode *curr = root;
          while (curr != nullptr) {
              if (curr->left != nullptr) {
                  auto next = curr->left;
                  auto predecessor = next;
                  while (predecessor->right != nullptr) {
                      predecessor = predecessor->right;
                  }
                  predecessor->right = curr->right;
                  curr->left = nullptr;
                  curr->right = next;
              }
              curr = curr->right;
          }
      }
  };

• 头插法：按照右子树 - 左子树 - 根的顺序 DFS 这棵树，DFS 的同时，记录当前链表的头节点为 head。一开始 head 是空节点，具体来说：

  • 如果当前节点为空，返回
  • 递归右子树
  • 递归左子树
  • 把 root.left 置为空
  • 头插法，把 root 插在 head 的前面，也就是 root.right=head
  • 现在 root 是链表的头节点，把 head 更新为 root

• 代码实现

  class Solution {
      TreeNode* head;
  public:
      void flatten(TreeNode* root) {
          if (root == nullptr) {
              return;
          }
          flatten(root->right);
          flatten(root->left);
          root->left = nullptr;
          root->right = head; // 头插法，相当于链表的 root->next = head
          head = root; // 现在链表头节点是 root
      }
  };

• 头插法需要用到一个在 DFS 之外的变量 head，能否只在 DFS 中解决？

• 假设已经计算出了左子树和右子树的链表，那么只需要将两者连起来，在这个过程中，DFS 需要返回链表的尾节点，来方便连接两个链表

  class Solution {
      TreeNode* dfs(TreeNode* root) {
          if (root == nullptr) {
              return nullptr;
          }
          TreeNode* left_tail = dfs(root->left);
          TreeNode* right_tail = dfs(root->right);
          if (left_tail) {
              left_tail->right = root->right; // 左子树链表 -> 右子树链表
              root->right = root->left; // 当前节点 -> 左右子树合并后的链表
              root->left = nullptr;
          }
          return right_tail ? right_tail : left_tail ? left_tail : root;
      }
  
  public:
      void flatten(TreeNode* root) {
          dfs(root);
      }
  };