All Nodes Distance K in Binary Tree

0863.All-Nodes-Distance-K-in-Binary-Tree/memo.md

@@ -0,0 +1,12 @@
+# 863. All Nodes Distance K in Binary Tree
+
+## step1
+とりあえず動くものを書いてから整えた。最初のものは変数名や関数が適当だが残しておく。
+DFS。
+
+## step2
+BFSも書いてみる。
+
+https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/solutions/143729/python-dfs-and-bfs-by-lee215-v8do/?envType=problem-list-v2&envId=7p5x763
+
+DFSとBFSの解法。

0863.All-Nodes-Distance-K-in-Binary-Tree/step1_dfs.py

@@ -0,0 +1,61 @@
+# Definition for a binary tree node.
+# class TreeNode:
+#     def __init__(self, x):
+#         self.val = x
+#         self.left = None
+#         self.right = None
+
+
+class Solution:
+    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> list[int]:
+        node_to_parent = {root: None}
+        root_of_target = {root: False}
+
+        def find_target(node):
+            root_of_target[node] = False
+            if node == target:
+                root_of_target[node] = True
+                parent = node_to_parent[node]
+                while parent is not None:
+                    root_of_target[parent] = True
+                    parent = node_to_parent[parent]
+            for child in (node.left, node.right):
+                if child is not None:
+                    node_to_parent[child] = node
+                    find_target(child)
+
+        find_target(root)
+
+        result = []
+
+        def distance_k_of_children(node, count):
+            if count == k:
+                result.append(node.val)
+                return
+
+            for child in (node.left, node.right):
+                if child is not None:
+                    distance_k_of_children(child, count + 1)
+
+        distance_k_of_children(target, 0)
+        seen = {target}
+
+        def distance_k_of_others(node, count):
+            if node in seen or node is None:
+                return
+            seen.add(node)
+            if count == k:
+                result.append(node.val)
+                return
+            count_children = count + 1
+            for child in (node.left, node.right):
+                if child is not None:
+                    distance_k_of_others(child, count + 1)
+            parent = node_to_parent[node]
+            if parent is not None:
+                count_parent = count + 1
+                distance_k_of_others(parent, count + 1)
+
+        distance_k_of_others(node_to_parent[target], 1)
+
+        return result

0863.All-Nodes-Distance-K-in-Binary-Tree/step1_dfs_revised.py

@@ -0,0 +1,38 @@
+# Definition for a binary tree node.
+class TreeNode:
+    def __init__(self, x):
+        self.val = x
+        self.left = None
+        self.right = None
+
+
+class Solution:
+    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> list[int]:
+        node_to_parent = {}
+
+        def find_parents(node, parent):
+            if node:
+                node_to_parent[node] = parent
+                find_parents(node.left, node)
+                find_parents(node.right, node)
+
+        find_parents(root, None)
+
+        result = []
+        seen = set()
+
+        def traverse(node, distance):
+            if not node or node in seen:
+                return
+            seen.add(node)
+
+            if distance == k:
+                result.append(node.val)
+                return
+
+            traverse(node.left, distance + 1)
+            traverse(node.right, distance + 1)
+            traverse(node_to_parent[node], distance + 1)
+
+        traverse(target, 0)
+        return result

0863.All-Nodes-Distance-K-in-Binary-Tree/step2_bfs.py

@@ -0,0 +1,43 @@
+import collections
+
+
+# Definition for a binary tree node.
+class TreeNode:
+    def __init__(self, x):
+        self.val = x
+        self.left = None
+        self.right = None
+
+
+class Solution:
+    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> list[int]:
+        if k == 0:
+            return [target.val]
+
+        graph = collections.defaultdict(list)
+
+        def create_graph(node, parent):
+            for child in (node.left, node.right):
+                if child is not None:
+                    graph[node].append(child)
+                    create_graph(child, node)
+            if parent is not None:
+                graph[node].append(parent)
+
+        create_graph(root, None)
+
+        seen = {target}
+        frontier = [target]
+        distance = 0
+        while frontier:
+            next_frontier = []
+            for node in frontier:
+                seen.add(node)
+                for child in graph[node]:
+                    if child is not None and child not in seen:
+                        next_frontier.append(child)
+            frontier = next_frontier
+            distance += 1
+            if distance == k:
+                break
+        return [node.val for node in frontier]