【刷题】计算三叉搜索树的高度

# line1 = int(sys.stdin.readline())
# nums = list(map(int, sys.stdin.readline().split()))

# line1 = 3
# nums = [5000, 4000 ,3000]
line1 = 5
nums = [5000, 2000, 5000, 8000, 1800]


class Tree(object):
    def __init__(self, val, left=None, mid=None, right=None):
        self.val = val
        self.left: Tree = left
        self.mid: Tree = mid
        self.right: Tree = right

    def insert(self, num):
        if num < self.val - 500:
            if self.left:
                self.left.insert(num)
            else:
                self.left = Tree(num)
        elif num > self.val - 500:
            if self.right:
                self.right.insert(num)
            else:
                self.right = Tree(num)
        else:
            if self.mid:
                self.mid.insert(num)
            else:
                self.mid = Tree(num)

    def get_height(self):
        height = 1
        if self.left:
            height = max(height, self.left.get_height() + 1)
        if self.mid:
            height = max(height, self.mid.get_height() + 1)
        if self.right:
            height = max(height, self.right.get_height() + 1)
        return height


root = Tree(nums[0])

for num in nums[1:]:
    root.insert(num)

print(root.get_height())