leetcode之236二叉树的最近公共祖先Golang

题目描述

给定一个二叉树, 找到该树中两个指定节点的最近公共祖先。

百度百科中最近公共祖先的定义为：“对于有根树 T 的两个结点 p、q，最近公共祖先表示为一个结点 x，满足 x 是 p、q 的祖先且 x 的深度尽可能大（一个节点也可以是它自己的祖先）。”

例如，给定如下二叉树: root = [3,5,1,6,2,0,8,null,null,7,4]

示例 1:

输入: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 1
输出: 3
解释: 节点 5 和节点 1 的最近公共祖先是节点 3。

示例 2:

输入: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 4
输出: 5
解释: 节点 5 和节点 4 的最近公共祖先是节点 5。因为根据定义最近公共祖先节点可以为节点本身。

说明:

• 所有节点的值都是唯一的。
• p、q 为不同节点且均存在于给定的二叉树中。

算法描述

本题主要的思路就是采用二叉树的后序遍历

利用栈来实现二叉树的后序遍历

核心思想就是当遍历到一个结点，是第一次找到p,q中的一个，那么就将当前结点设置为当前的最近祖先lowest，然后继续遍历后面的结点，按照后序遍历的顺序，直接按照下面的顺序来吧

• 将根结点入栈
• 判断当前栈顶是否存在左子树，如果存在左子树，就将左子树的根结点入栈，然后进入下一次循环
  • 每遍历到一个结点，判断是否是p,q中的一个，如果不是就直接跳过
  • 如果是，分第一次找到和第二次站到
    • 第一次找到，那么就将当前结点设置为当前的最近祖先lowest
    • 如果是第二次找到，那么就直接返回lowest，结束算法
• 如果当前栈顶元素不存在左子树，但是存在右子树，将右子树的根节点入栈
  • 每遍历到一个结点，判断是否是p,q中的一个，如果不是就直接跳过
  • 如果是，分第一次找到和第二次站到
    • 第一次找到，那么就将当前结点设置为当前的最近祖先lowest
    • 如果是第二次找到，那么就直接返回lowest，结束算法
• 如果当前栈顶元素既不存在左子树，也不存在右子树，那么就弹出当前的栈顶元素
  • 如果弹出的元素是我们前面设置的lowest，那么就将lowest重新设置为栈顶的下一个元素
  • 判断当前弹出的元素是栈顶下一个元素的左孩子还是右孩子
    • 左孩子：
      • 按照后序遍历的顺序，接下来该遍历右子树，判断右子树是否存在，如果存在就将右子树的根节点入栈
        • 检查当前遍历到的元素是否是p,q中的一个，然后做出和上面检查一样的处理
        • 然后退出到最外层的循环（就是判断是否存在左孩子，是否存在右孩子的循环）
    • 右孩子：
      • 直接继续重复当前将栈顶元素弹出栈的循环

总结，其实就是利用栈实现二叉树的后序遍历，然后遍历的时候检查遍历到的元素是否是p,q中的一个，如果是第一次遇到p,q，那么就设置开始设置最近祖先，然后第二次遇到p,q，就直接返回公共祖先。在中间将元素弹出栈的时候，判断弹出栈的元素是不是当前设置的公共祖先，如果是的话，就将公共祖先设置为栈顶的下一个元素。

代码

func lowestCommonAncestor(root, p, q *TreeNode) *TreeNode {
	if root == nil {
		return nil
	}
	nodeStack := []*TreeNode{root}
	findFirst := false
	findSecond := false
	var lowest *TreeNode
	// 首先将根结点入栈
	// 判断根结点是否是两个结点之一，如果是的话，就相当于找到了第一个结点，
	// 当前的祖先公共祖先就是根节点
	if root == p || root == q {
		findFirst = true
		lowest = root
	}
	for len(nodeStack) != 0 {
		// 第一层循环，将结点入栈
		if nodeStack[len(nodeStack)-1].Left == nil && nodeStack[len(nodeStack)-1].Right == nil {
			// 栈顶结点既没有左子树，又没有右子树，那么就pop栈顶结点
			for {
				// 进入内层循环，pop栈顶结点
				prePop := nodeStack[len(nodeStack)-1]
				if prePop == lowest {
					// 如果pop出的栈顶结点是前面设置的最近祖先，那么就将最近祖先设置为它的父节点，也就是下一个栈顶元素
					lowest = nodeStack[len(nodeStack)-2]
				}
				nodeStack = nodeStack[:len(nodeStack)-1]
				if nodeStack[len(nodeStack)-1].Left == prePop {
					// 被pop出的结点是下一个栈顶元素的左孩子
					// 如果是左孩子，那么按照后序遍历的顺序，还要检查它的右子树
					// 当右子树不为空，那么就将右子树的根结点入栈
					if nodeStack[len(nodeStack)-1].Right != nil {
						if nodeStack[len(nodeStack)-1].Right == p || nodeStack[len(nodeStack)-1].Right == q {
							// 检查当前查找到的结点是否是两个结点之一，如果是的话，判断第几次找到
							if !findFirst {
								findFirst = true
								lowest = nodeStack[len(nodeStack)-1].Right
							} else {
								// 找到了第二个结点，结束当前层循环
								findSecond = true
								break
							}
						}
						// 右子树的根结点入栈
						nodeStack = append(nodeStack, nodeStack[len(nodeStack)-1].Right)
						break
					} else {
						// 不存在右子树，那么继续当前循环，也就是继续将栈顶元素出栈
						continue
					}
				} else {
					// 被pop出的元素是下一个栈顶元素的右孩子，根据后序遍历的顺序，继续pop父节点
					// 所以继续当前层循环
					continue
				}
			}
			if findSecond {
				break
			}
		} else if nodeStack[len(nodeStack)-1].Left != nil {
			// 如果栈顶结点存在左子树，那么将左子树的根节点入栈
			nodeStack = append(nodeStack, nodeStack[len(nodeStack)-1].Left)
			// 检查当前入栈的结点是不是两个结点之一
			if nodeStack[len(nodeStack)-1] == p || nodeStack[len(nodeStack)-1] == q {
				// 第一次找到两个结点之一，那么设定当前栈顶结点为祖先
				if !findFirst {
					findFirst = true
					lowest = nodeStack[len(nodeStack)-1]
				} else {
					// 如果是第二次找到了两个结点之一，那么就直接退出循环，返回祖先
					break
				}
			}
		} else {
			// 栈顶结点的左子树为空，又子树不为空，那么就将右子树的根节点入栈
			nodeStack = append(nodeStack, nodeStack[len(nodeStack)-1].Right)
			if nodeStack[len(nodeStack)-1] == p || nodeStack[len(nodeStack)-1] == q {
				// 检查是否找到两个结点之一，并且是第几次找到，判断和上面的类似
				if !findFirst {
					findFirst = true
					lowest = nodeStack[len(nodeStack)-1]
				} else {
					break
				}
			}
		}
	}
	return lowest
}