cousions_in_binary_tree.py

# Definition for a  binary tree node
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

from collections import deque
class Solution:
    # @param A : root node of tree
    # @param B : integer
    # @return a list of integers
    def solve(self, A, B):
        queue = deque()
        queue.append([A, A])
        target_parent = None
        while queue:
            # current will an list with current node address as first item and its parent is second item
            current = queue.popleft()
            current_node = current[0]
            parent = current[1]
            if current_node.val == B:
                # storing the parent of target node so that we can exclude it from cousions
                target_parent = parent
                break
            if current_node.left:
                queue.append([current_node.left, current_node])
            if current_node.right:
                queue.append([current_node.right, current_node])
        
        cousins = []
        while queue:
            current = queue.popleft()
            current_node = current[0]
            parent = current[1]
            # if node has same parent as target node so we are sipping it
            if parent != target_parent:
                cousins.append(current_node.val)
        return cousins