2021-06-04：给定三个参数：二叉树的头节点head，树上某个节点target，正数K，从target开始，可以向上走或者

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)
    }
}