Flatten Binary Tree

FlattenBinaryTree.py

# Leetcode: 114. Flatten Binary Tree to Linked List

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

    def flattenBT(self, node: Optional[TreeNode]) -> Optional[TreeNode]:
        if node:
            # get left and right of node
            node_lft, node_rgt = node.left, node.right

            # disconnect left and right of node
            node.left = None
            node.right = None

            # connect the right of node to left node if exist else to right node
            node.right = node_lft if node_lft else node_rgt

            # flatten left of node
            lft_tail = self.flattenBT(node_lft)

            # connect the left subtree to right node if it exist
            if lft_tail:
                lft_tail.right = node_rgt

            # flatten right of node
            rgt_tail = self.flattenBT(node_rgt)

            return rgt_tail or lft_tail or node

    def flatten(self, root: Optional[TreeNode]) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        self.flattenBT(root)
        return