1382. Balance a Binary Search Tree

Information

  • Diffculty: Medium

  • Created: 2020-03-15 11:55:48

  • Last Motified: 2020-03-15 11:55:48

  • Tags: Binary Search Tree

Intuition

This problem could be solved in following steps,

  1. Get the number of all nodes.

  2. With the number of all nodes, we could easily get the depth of the tree and the number of nodes in the last row.

  3. We would need a function to get the inorder of the tree with nodes one by one. The easier way is to save the inorder of the original tree into a list. But we could use a non-recursive inorder to do it with space complexity O(lg(n)).

  4. Construct the final tree with the depth and number of nodes in the last tree.

Solution

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

class Solution:
    def balanceBST(self, root: TreeNode) -> TreeNode:
        if root is None:
            return None

        # Count the overall nodes
        nodes = self._dfs_count(root)

        # Get the depth of the tree
        partial_depth = 1
        while nodes >= ((2 ** partial_depth) - 1):
            partial_depth += 1
        full_depth = partial_depth - 1
        partial_nodes = nodes - ((2 ** full_depth) - 1)

        # Initial inorder traversal stack
        # The inorder interate function is _inorder_next
        self._inorder_stack = [[root, 0]]

        self.last_row = partial_nodes
        return self._inorder_construct(full_depth)

    def _dfs_count(self, node):
        if node is None:
            return 0
        return 1 + self._dfs_count(node.left) + self._dfs_count(node.right)

    def _inorder_next(self):
        while self._inorder_stack:
            node, state = self._inorder_stack.pop()
            if state == 0:
                self._inorder_stack.append([node, 1])
                if node.left:
                    self._inorder_stack.append([node.left, 0])
            else:
                if node.right:
                    self._inorder_stack.append([node.right, 0])
                return node.val

        return None

    def _inorder_construct(self, depth):
        if depth == 0:
            if self.last_row == 0:
                return None
            self.last_row -= 1
            return TreeNode(self._inorder_next())
        node = TreeNode(None)
        node.left = self._inorder_construct(depth - 1)
        node.val = self._inorder_next()
        node.right = self._inorder_construct(depth - 1)
        return node