二叉树两节点的最短距离(一次dfs完成)

最新推荐文章于 2022-04-01 20:48:40 发布

原创最新推荐文章于 2022-04-01 20:48:40 发布 · 806 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#深度优先 #leetcode #算法

字节跳动面试题目，通过深度优先搜索（DFS）解决二叉树中寻找两个给定节点最短距离的问题。与LeetCode 236题类似但有区别。在DFS过程中，利用返回值表示节点是否被访问，记录节点间的高度差，以确定节点间的最短路径。具体实现包括多种情况的判断，如当前节点是目标之一、子树中存在另一个目标节点等，最终找到最短距离。

字节面试问到的，leetcode上没看到，记录一下自己的做法

此题近似于leetcode 236.二叉树的最近公共祖先，但有较大不同

做法为：对二叉树进行dfs遍历，递归函数写法为func dfs(root,p,q *TreeNode,ans *int) int，

其中root为当前遍历的节点，p,q为要求的两个节点，ans是答案的指针，返回int值

总体思想为利用返回值表征两个节点中的一个是否出现过，并记录距当层的高度，若根节点、左子树、右子树中的二者同时出现p,q则更新答案

分多种情况讨论：

root==nil，遍历到空节点，返回0
root!=nil时，分情况讨论
1. 先遍历左右子节点，返回值记为left,right
2. 若当前root是p,q中的一个
  1. 判断左右子树中是否有另一个节点(left!=0或right!=0)，如果有，则直接更新答案，返回
  2. 否则返回1，告诉上层节点自己出现过并记录相差的高度
3. root不是p,q中的一个
  1. 若left,right均不为0，代表左右子树都出现过p,q，则更新答案为left+right，返回
  2. 若均为0，则p,q不在当前树中，返回0
  3. 剩下的情况是left,right中有1个不为0且自己(root)不是另一个，则向上传递left+1或right+1(非0的那个+1)记录相差的高度，这样一直向上传递到p,q中另一个节点或者传到p,q的最近公共祖先就会更新答案

可见只有两种情况会更新答案

p,q中一个节点是另一个节点的父节点，假定p是q父节点，p发现q出现在自己子树中，更新答案为p,q二者相差的高度
遍历到了p,q的最近公共祖先，记为A，A发现p,q在左右子树都出现过，则更新答案为二者高度的和（注意A是唯一的，只会更新一次）

go代码：

package main

import "fmt"

type TreeNode struct {
	val         int
	Left, Right *TreeNode
}

func dfs(root, p, q *TreeNode, ans *int) int {
	if root == nil {
		return 0
	}
	left := dfs(root.Left, p, q, ans)
	right := dfs(root.Right, p, q, ans)
	if root == p || root == q {
		if left != 0 {
			*ans = left
			return 0
		}
		if right != 0 {
			*ans = right
			return 0
		}
		return 1
	}
	if left != 0 && right != 0 {
		*ans = left + right
		return 0
	}
	if left == 0 && right == 0 {
		return 0
	}
	if left != 0 {
		return left + 1
	}
	return right + 1
}

func main() {
	A := &TreeNode{1, nil, nil}
	B := &TreeNode{2, nil, nil}
	C := &TreeNode{3, nil, nil}
	D := &TreeNode{4, nil, nil}
	E := &TreeNode{5, nil, nil}
	F := &TreeNode{6, nil, nil}
	A.Left, A.Right = B, C
	B.Left, B.Right = D, E
	D.Left = F
	ans := 0
	dfs(A, F, C, &ans)
	fmt.Println(ans) //4
	dfs(A, F, B, &ans)
	fmt.Println(ans) //2
	dfs(A, F, E, &ans)
	fmt.Println(ans) //3
	dfs(A, A, F, &ans)
	fmt.Println(ans) //3
	dfs(A, A, B, &ans)
	fmt.Println(ans) //1
}