Advertisement
kai-rocket

Cousins in Binary Tree

Mar 26th, 2021
578
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.92 KB | None | 0 0
  1. from collections import deque
  2.  
  3. # Definition for a binary tree node.
  4. # class TreeNode:
  5. #     def __init__(self, val=0, left=None, right=None):
  6. #         self.val = val
  7. #         self.left = left
  8. #         self.right = right
  9. class Solution:
  10.     def isCousins(self, root: TreeNode, x: int, y: int) -> bool:
  11.         # Each element in my queue is a dict with attributes curr_node, parent_val, and depth
  12.         queue = deque([{'curr_node': root, 'parent_val': None, 'depth': 0}])
  13.        
  14.         # Iterate BFS until we exhaust the tree
  15.         while len(queue) > 0:
  16.             # Access the current node
  17.             curr_node_data = queue.popleft()
  18.             curr_node = curr_node_data['curr_node']
  19.            
  20.             # If I find x or y:
  21.             if curr_node.val == x or curr_node.val == y:
  22.                 # Search remainder of queue for the other value
  23.                 for node_data in queue:
  24.                     # If other value exists in queue:
  25.                     if (curr_node.val == x and node_data['curr_node'].val == y) or (curr_node.val == y and node_data['curr_node'].val == x):
  26.                         # Return True if x and y nodes have different parents and same level, False otherwise    
  27.                         return curr_node_data['parent_val'] != node_data['parent_val'] and curr_node_data['depth'] == node_data['depth']
  28.                 # Return false if X or Y is found but the "opposite" is not found in the current queue
  29.                 return False
  30.            
  31.             # Continue iteration
  32.             if curr_node.left:
  33.                 queue.append({'curr_node': curr_node.left, 'parent_val': curr_node.val, 'depth': curr_node_data['depth'] + 1})
  34.             if curr_node.right:
  35.                 queue.append({'curr_node': curr_node.right, 'parent_val': curr_node.val, 'depth': curr_node_data['depth'] + 1})
  36.        
  37.         # If x and y do not exist in tree, return False
  38.         return False
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement