给定三个参数：二叉树的头节点head，树上某个节点target，正数K，从target开始，可以向上走或者向下走。返回与target的距离是K的所有节点。
记录父节点的map，key是当前节点，value是父节点。
访问集合，凡是节点被访问过，放在这个集合中。
队列，广度优先遍历。
代码用golang编写。代码如下：
package main
import "fmt"
func main() {
root := &Node{Val: 1}
root.Left = &Node{Val: 2}
root.Right = &Node{Val: 3}
root.Right.Right = &Node{Val: 6}
root.Left.Left = &Node{Val: 4}
root.Left.Right = &Node{Val: 5}
root.Left.Right.Left = &Node{Val: 7}
root.Left.Right.Right = &Node{Val: 8}
target := root.Left
ret := distanceKNodes(root, target, 2)
for i := 0; i < len(ret); i++ {
fmt.Println(ret[i].Val)
}
}
type Node struct {
Val int
Left *Node
Right *Node
}
func distanceKNodes(root *Node, target *Node, K int) []*Node {
parents := make(map[*Node]*Node)
parents[root] = nil
createParentMap(root, parents)
queue := make([]*Node, 0)
visited := make(map[*Node]struct{})
queue = append(queue, target)
visited[target] = struct{}{}
curLevel := 0
ans := make([]*Node, 0)
for len(queue) > 0 {
size := len(queue)
for size > 0 {
size--
cur := queue[0]
queue = queue[1:]
if curLevel == K {
ans = append(ans, cur)
}
if cur.Left != nil {
if _, ok := visited[cur.Left]; !ok {
visited[cur.Left] = struct{}{}
queue = append(queue, cur.Left)
}
}
if cur.Right != nil {
if _, ok := visited[cur.Right]; !ok {
visited[cur.Right] = struct{}{}
queue = append(queue, cur.Right)
}
}
if parents[cur] != nil {
if _, ok := visited[parents[cur]]; !ok {
visited[parents[cur]] = struct{}{}
queue = append(queue, parents[cur])
}
}
}
curLevel++
if curLevel > K {
break
}
}
return ans
}
func createParentMap(cur *Node, parents map[*Node]*Node) {
if cur == nil {
return
}
if cur.Left != nil {
parents[cur.Left] = cur
createParentMap(cur.Left, parents)
}
if cur.Right != nil {
parents[cur.Right] = cur
createParentMap(cur.Right, parents)
}
}执行结果如下：
