Leetcode 863 All Nodes Distance K in Binary Tree

发表于 |



这道题思想非常简单，但是可能我代码写得有点复杂（因为我回溯用的不好），先着重看思路把，代码实现方法尽量讲一讲

分析：

  1. 我们先考虑最简单的情况，如果根结点就是目标节点，那么找离根结点距离为K的结点很容易，使用bfs找到第k层所有结点即可，现在定义getAns(root,K)来处理这种情况！
  2. 如下图，如果我们找到目标结点设为self.root，那么圆圈所圈部分中的结果，可直接用getAns(self.root,K)得到
  3. 那么圆圈外的部分，我们应该怎么得到呢？如果我们能想办法将圆圈外的部分也化为以目标节点为根节点的形式就直接解决了， 这次看一个复杂的情形，我们需要把圆圈内所有箭头（即目标节点至根结点路径中所有箭头）反转，才能构建出我们想要的形式，构建完之后使用getAns(self.root,K)计算结果。

思路：
着重讲下search()函数

  1. 当我们找到目标节点后，先计算一次结果，如分析中第2步。然后将其左右孩子全部抹去，将其本身父结点设为其孩子节点。
  2. 而为了尽量不改变父节点本身的结构，如左孩子为目标节点，则将父结点的left先指向None,方便之后再指向父节点的父节点，依次类推。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# 别看代码挺长的，其实search()函数里面做的事是重复的，只是分了左右节点而已，思路还是比较明朗

import Queue
class Solution(object):
    def __init__(self):
        self.res = []
        self.root = None
        self.visited = set()
    
    def distanceK(self, root, target, K):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type K: int
        :rtype: List[int]
        """
        if root is target: return self.getAns(root, K)
        self.search(root, target,K)
        self.res += self.getAns(self.root, K)
        return list(set(self.res))
    
    # 当根节点为目标节点时，使用bfs找到第k层的所有节点
    def getAns(self,root,K):
        Q = Queue.deque([(root,0)])
        res = []
        while Q:
            node,level = Q.popleft()
            if level == K:
                res.append(node.val)
            if node.left:
                Q.append((node.left,level+1))
            if node.right:
                Q.append((node.right,level+1))
        return res
    
    # 找到目标节点并设为根节点，然后转换树结构
    def search(self,root,target,K):
        if not root: return None
        l = self.search(root.left,target,K)
        r = self.search(root.right, target,K)
        # 找到目标节点，先进行一次计算，然后转换树结构
        # 将其父节点放入visited，表明这是待转换路径中的节点
        if l is target:
            self.root = l
            self.res = self.getAns(l,K)
            l.left = None
            l.right = root
            root.left = None
            self.visited.add(root)
        if r is target:
            self.root = r
            self.res = self.getAns(r,K)
            r.right = None
            r.left = root
            root.right = None
            self.visited.add(root)

        # 找到待转换路径中的节点，转换树结构，再将父节点加入visited，依次类推
        # 如思路中第2点所说，尽量不改变父节点本身结构，所有需判断一下
        if l in self.visited:
            if l.left == None:
                l.left = root
            elif l.right == None:
                l.right = root
            root.left = None
            self.visited.add(root)

        if r in self.visited:
            if r.left == None:
                r.left = root
            elif r.right == None:
                r.right = root
            root.right = None
            self.visited.add(root)
        return root


57 / 57 test cases passed.
difficulty: medium
Runtime: 50 ms