Leetcode 863. All Nodes Distance K in Binary Tree

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

All Nodes Distance K in Binary Tree - LeetCode

문제

이진트리의 root와 함께 목표 노드의 노드 정보가 담긴 target 그리고 정수 k가 주어진다. 이때 거리 target 노드에서 거리 k만큼 떨어진 노드들의 값을 모두 반환하라

 

답의 순서는 상관없다

 

코드

from collections import defaultdict

class Solution:

    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:

        # find target node

        m = self.makeMap(root)

        # print(m)

        q = [target.val]

        visited = [target.val]

        while k > 0:

            temp_q = []

            while q:

                for node in m[q.pop()]:

                    if node not in visited:

                        temp_q += [node]

                    visited += [node]

            q = temp_q

            k -= 1

        return q

    def makeMap(self, root):

        m = defaultdict(list)

        q = [root]

        while q:

            cursor = q.pop()

            if cursor.left is not None:

                m[cursor.val].append(cursor.left.val)

                m[cursor.left.val] += [cursor.val]

                q.append(cursor.left)

            if cursor.right is not None:

                m[cursor.val].append(cursor.right.val)

                m[cursor.right.val] += [cursor.val]

                q.append(cursor.right)

        return m

---

시간은 좀 부족했지만..! 메모리 부분에서 뛰어난 성능을 보여줬으니 만-족

코드는 모든 노드에서 갈 수 있는 노드들에 대한 정보를 m에 담고 방문하지 않은 노드들을 계속 타고가는 형식으로 코드를 구현했다.