常数空间Invert Binary Tree与仿Morris法后序遍历

今天在地铁上浪费了好多时间……呜呜……做其他事都没有效率,就利用这些时间写字了。

前几天LeetCode的Invert Binary Tree火了。名称有点糟,无法准确描述要求。

主要思路是遍历二叉树,把访问操作修改为交换左右孩子,每个节点都交换一次即可。

如果采用前序或中序遍历,则子树交换发生在遍历某棵子树之前,会引起麻烦。因此我想到了后序遍历,在遍历完子树后再交换左右孩子,重心就是如何实现常数空间复杂度的后序遍历。某些资料/问题称之为Morris post-order traversal。

考虑树遍历的递归算法,一个节点的访问次数为度(孩子数)+1。对于前序遍历,第一次访问产生输出;对于中序遍历,第二次访问产生输出。鉴于输出节点信息的那次访问比较简单,把它和下一次访问合并。

因此前序遍历最多只有两次访问,访问时执行的操作如下:

  • 若左孩子非空则遍历左子树
  • 访问当前节点。遍历右子树

中序遍历相仿:

  • 访问当前节点。若左孩子非空则遍历左子树
  • 遍历右子树

代码如下:

struct Node { Node *l, *r; void visit(); };

while (p) {
  Node *q = p->l;
  if (q) {
    while (q->r && q->r != p) q = q->r;
    if (q->r == p)
      q->r = NULL;
    else {
      if (t == PREORDER)
        p->visit();
      q->r = p;
      p = p->l;
      continue;
    }
  } else if (t == PREORDER)
    p->visit();
  if (t == INORDER)
    p->visit();
  p = p->r;
}

定义一个概念,节点x的右链表示代码while (x) x = x->r;中x指向过的所有节点,最右节点为最后一个不等于NULL(nullptr)的x指针值。

遍历完某棵子树后需要回到中序序列的后继。一个左孩子非空的节点会成为当前访问的节点(p)两次。第一次p左子树的最右节点q尚未建立线索,令其右孩子指向中序序列后继p(线索);第二次q的线索存在,指向p,删除之。

对于后序遍历,一个度为2的节点只有在第3次访问时才输出,而Morris in-order traversal算法中一个节点最多只会作为p两次,无法提供3次访问。我们需要挖掘更多的信息。实际上除了根的右链以外的节点都被指向过3到4次:

  • 该节点所在右链的父节点为p时,被变量q指向
  • 第一次被p指向。此时左孩子为空或左子树最右节点的thread不存在
  • 若左孩子非空,会第二次被p指向,左子树右链的thread存在
  • 该节点所在右链的父节点为p时,被变量q指向

如果把第2次变量q指向当作第3次访问,则除了根的右链以外的节点都能被访问2(左孩子为空)或3(左孩子非空)次,如果创建一个虚拟的根,原来的根作为它的左孩子,则根的右链的情况也和其他节点相同了,可以统一处理。

while (q->r && q->r != p) q = q->r;中若倒序输出各个q的值,则得到了后序序列。常数空间复杂度实现倒序输出时需要用到链表翻转的技巧,翻转右链(可以常数空间实现)再沿着链输出,之后再翻转回来(还原)。

// morris_postorder_traversal的辅助函数
void reverse_right_chain(Node *x, Node *y)
{
  Node *p = x, *q = x->r, *r;
  while (p != y) {
    r = q->r;
    q->r = p;
    p = q;
    q = r;
  }
}

void morris_postorder_traversal(Node *p)
{
  Node aux;
  aux.l = p;
  aux.r = NULL;
  p = &aux;
  while (p) {
    Node *q = p->l;
    if (q) {
      while (q->r && q->r != p) q = q->r;
      if (q->r == p) {
        reverse_right_chain(p->l, q);
        for (Node *r = q; ; r = r->r) {
          r->visit();
          if (r == p->l) break;
        }
        reverse_right_chain(q, p->l);
        q->r = NULL;
      } else {
        q->r = p;
        p = p->l;
        continue;
      }
    }
    p = p->r;
  }
}

回到Invert Binary Tree,其实不要求严格的遍历顺序,因此我们不必翻转两次右链达到后序效果:

class Solution {
public:
  TreeNode* invertTree(TreeNode* root) {
    TreeNode aux(0), *p = &aux;
    aux.left = root;
    while (p) {
      TreeNode *q = p->left;
      if (q) {
        while (q->right && q->right != p) q = q->right;
        if (q->right == p) {
          for (TreeNode *r = p->left; ; r = r->left) {
            swap(r->left, r->right);
            if (r == q) break;
          }
          q->left = NULL;
        } else {
          q->right = p;
          p = p->left;
          continue;
        }
      }
      p = p->right;
    }
    return aux.left;
  }
};

其他题可以参考LeetCode solutions,持续更新中。