Leetcode 893 Solution

This article provides solution to leetcode question 893 (all-nodes-distance-k-in-binary-tree)
Categories:
Tags:
Access this page by simply typing in "lcs 893" in your browser address bar if you have bunnylol configured.

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree

Solution

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def dfs(self, root, par):
        if not root:
            return

        root.par = par
        self.dfs(root.left, root)
        self.dfs(root.right, root)

    def distanceK(self, root, target, K):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type K: int
        :rtype: List[int]
        """
        self.dfs(root, None)

        q = [target]
        seen = {target}
        d = 0

        while q:
            if d == K:
                return [node.val for node in q]

            s = len(q)
            for _ in range(s):
                node = q.pop(0)
                for nei in [node.par, node.left, node.right]:
                    if not nei or nei in seen:
                        continue
                    seen.add(nei)
                    q.append(nei)

            d += 1

        return []