fork download 
  1. # author: Eldan
  2. # task: leetcode.com/problems/all-nodes-distance-k-in-binary-tree/
  3. from collections import deque
  4.  
  5.  
  6. class TreeNode:
  7.     def __init__(self, val):
  8.         self.val = val
  9.         self.left = None
  10.         self.right = None
  11.  
  12.  
  13. class Solution:
  14.     def _directed_parent_connection(self, node, parent=None):
  15.         if node:
  16.             node.parent = parent
  17.             self._directed_parent_connection(node.left, node)
  18.             self._directed_parent_connection(node.right, node)
  19.  
  20.     def k_distance_nodes(self, root, target, k) -> []:
  21.         if not root:
  22.             return []
  23.  
  24.         self._directed_parent_connection(root)
  25.  
  26.         queue = deque([(target, 0)])
  27.         visited_nodes = {target}
  28.  
  29.         while queue:
  30.             if queue[0][1] == k:
  31.                 return [node.val for node, distance in queue]
  32.  
  33.             node, distance = queue.popleft()
  34.             for neighbor in (node.left, node.right, node.parent):
  35.                 if neighbor and neighbor not in visited_nodes:
  36.                     queue.append((neighbor, distance + 1))
  37.                     visited_nodes.add(neighbor)
  38.  
  39.         return []
  40.  
  41.  
  42. if __name__ == '__main__':
  43.     tree = TreeNode(3)
  44.     tree.left = TreeNode(5)
  45.     tree.right = TreeNode(1)
  46.  
  47.     tree.left.left = TreeNode(6)
  48.     tree.left.right = TreeNode(2)
  49.     tree.right.left = TreeNode(0)
  50.     tree.right.right = TreeNode(8)
  51.  
  52.     tree.left.right.left = TreeNode(7)
  53.     tree.left.right.right = TreeNode(4)
  54.  
  55.     """
  56.                 3
  57.               /   \
  58.            [5]		1
  59.           /  \    /  \
  60.         6     2  0	  8
  61.             /  \
  62.            7    4	
  63.     """
  64.     target = tree.left
  65.     k = 2
  66.  
  67.     solution = Solution()
  68.     result = solution.k_distance_nodes(tree, target, k)
  69.  
  70.     if not result:
  71.         print('There is no k distinated nodes in the tree')
  72.     else:
  73.         print(result)
  74.  
Success #stdin #stdout 0.02s 27712KB
comments (?)
stdin 
Standard input is empty
stdout 
[7, 4, 1]
https://ideone.com/GbrQQE
language:
Python 3 (python  3.12)
created:
7 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!