[Leetcode] 863. All Nodes Distance K in Binary Tree

whitehousechef·2025년 3월 10일

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/description/?envType=problem-list-v2&envId=7p5x763&sorting=W3sic29ydE9yZGVyIjoiREVTQ0VORElORyIsIm9yZGVyQnkiOiJGUkVRVUVOQ1kifV0%3D&page=1

initial

In a treenode, there are attributes of only left and right child nodes. But here you need to traverse to its root too so how?

If it were a regular graph question, I would have thought of making a separate graph of bidirectional graph but I was limiting myself.

my initial wrong attempt was trying to use 2 functions - 1 dfs for finding distance==k and 1 difs for finding the target value.

def distanceK(root: TreeNode, target, k):
    ans = []
    count = 0

    def dfs(node, dist):
        if dist == k:
            ans.append(node)
            return
        if not node:
            return
        dfs(node.left, dist + 1)
        dfs(node.right, dist + 1)

    def first(node, check):
        if node == target:
            return check
        if not node:
            return
        val1 = first(node.left, check + 1)
        if val1:
            dfs(node, val1)
        val2 = first(node.right, check + 1)
        if val2:
            dfs(node, val2)

    first(root, 0)

solution

we can just make a separate graph of parent <-> child. One very impt thing is this mkaing graph function. Always consider from the node POV, NOT node.left or right child's pov. Because it is simpler to implement if node: to check if it isnt null, rather than checking the children if they are null.

from collections import defaultdict,deque
class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int):
        graph = defaultdict(list)

        # always on the node POV, not node left or right
        def make_graph(node, parent=None):
            if node:
                if parent:
                    graph[parent.val].append(node.val)
                    graph[node.val].append(parent.val)
                make_graph(node.left, node)
                make_graph(node.right, node)

        make_graph(root)
        queue = deque()
        queue.append([target.val, 0])
        ans = []
        checker = set()
        checker.add(target.val)
        while queue:
            hola, count = queue.popleft()
            if count == k:
                ans.append(hola)
            for next in graph[hola]:
                if next not in checker:
                    checker.add(next)
                    queue.append((next, count + 1))
        return ans

complexity

o(n) time and space

0개의 댓글