Balance a Binary Search Tree LeetCode solution in Python

1
# Definition for a binary tree node.
2
# class TreeNode(object):
3
#     def __init__(self, val=0, left=None, right=None):
4
#         self.val = val
5
#         self.left = left
6
#         self.right = right
7
class Solution(object):
8
    def balanceBST(self, root):
9
        """
10
        :type root: TreeNode
11
        :rtype: TreeNode
12
        """
13
        arr = []
14
        self.flatTree(root, arr)
15
        return self.createTree(arr, 0, len(arr))
16

17
    def flatTree(self, root, arr):
18
        if not root:
19
            return
20
        self.flatTree(root.left, arr)
21
        arr.append(root)
22
        self.flatTree(root.right, arr)
23

24
    def createTree(self, arr, start, length):
25
        if length == 0:
26
            return None
27
        root = arr[start + length / 2]
28
        root.left = self.createTree(arr, start, length / 2)
29
        root.right = self.createTree(
30
            arr, start + length / 2 + 1, length - length / 2 - 1
31
        )
32
        return root