# 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
```