计算三叉搜索树的高度

题目描述

定义构造三叉搜索树规则如下:

每个节点都存有一个数, 当插入一个新的数时, 从根节点向下寻找, 直到找到一个合适的空节点插入.

查找的规则是:

• 如果数小于节点的数减去500, 则将数插入节点的左子树
• 如果数大于节点的数加上500, 则将数插入节点的右子树
• 否则, 将数插入节点的中子树

给你一系列数, 请按以上规则, 按顺序将数插入树中, 构建出一棵三叉搜索树, 最后输出树的高度.

输入描述

第一行为一个数 N, 表示有 N 个数, 1 ≤ N ≤ 10000

第二行为 N 个空格分隔的整数, 每个数的范围为 [1,10000]

输出描述

输出树的高度 (根节点的高度为1)

示例1

输入:

5
5000 2000 5000 8000 1800

输出:

3

说明:

示例2

输入:

3
5000 4000 3000

输出:

3

说明:

示例3

输入:

9
5000 2000 5000 8000 1800 7500 4500 1400 8100

输出:

4

说明:

题解

Python

class TreeNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.middle = None
        self.right = None

    # 向当前节点的子节点插入新的值
    # 如果当前节点的子节点不存在, 就创建它
    # 如果对应的子节点已经存在, 就在子节点中完成插入操作
    def insert(self, value):
        if self.value - value > 500:
            # 去左子节点
            if self.left:
                self.left.insert(value)
            else:
                self.left = TreeNode(value)
        elif value - self.value > 500:
            # 去右子节点
            if self.right:
                self.right.insert(value)
            else:
                self.right = TreeNode(value)
        else:
            # 去中间子节点
            if self.middle:
                self.middle.insert(value)
            else:
                self.middle = TreeNode(value)

# 构造三叉搜索树
def build_tree(array: list[int]) -> TreeNode:
    # 创建根节点
    root = TreeNode(array[0])
    for value in array[1:]:
        root.insert(value)
    return root

# 递归访问所有节点
def tree_height(root: TreeNode) -> int:
    if not root:
        return 0
    left_height = tree_height(root.left)
    middle_height = tree_height(root.middle)
    right_height = tree_height(root.right)
    child_height = max(left_height, middle_height, right_height)
    return 1 + child_height

def main():
    # 读取输入
    # 读取节点个数
    num_nodes = int(input())
    assert 1 <= num_nodes <= 10000
    # 读取所有节点的值, 存放到数组
    array = list(map(int, input().split()))
    assert len(array) == num_nodes

    # 构造三叉树
    root = build_tree(array)

    # 获取树的最大高度
    height = tree_height(root)

    # 打印结果
    print(height)

if __name__ == "__main__":
    main()

Rust

use std::io::{stdin, BufRead};

pub struct TreeNode {
    value: i32,
    left: Option<Box<Self>>,
    middle: Option<Box<Self>>,
    right: Option<Box<Self>>,
}

impl TreeNode {
    pub fn new(value: i32) -> Self {
        Self {
            value,
            left: None,
            middle: None,
            right: None,
        }
    }

    fn boxed_new(value: i32) -> Option<Box<Self>> {
        Some(Box::new(Self {
            value,
            left: None,
            middle: None,
            right: None,
        }))
    }

    // 向当前节点的子节点插入新的值
    // 如果当前节点的子节点不存在, 就创建它
    // 如果对应的子节点已经存在, 就在子节点中完成插入操作
    pub fn insert(&mut self, value: i32) {
        if self.value - value > 500 {
            // 去左子节点
            if let Some(left) = &mut self.left {
                left.insert(value);
            } else {
                self.left = Self::boxed_new(value);
            }
        } else if value - self.value > 500 {
            // 去右子节点
            if let Some(right) = &mut self.right {
                right.insert(value);
            } else {
                self.right = TreeNode::boxed_new(value);
            }
        } else {
            // 去中间子节点
            if let Some(middle) = &mut self.middle {
                middle.insert(value);
            } else {
                self.middle = TreeNode::boxed_new(value);
            }
        }
    }
}

// 构造三叉搜索树
fn build_tree(array: &[i32]) -> TreeNode {
    assert!(!array.is_empty());

    // 创建根节点
    let mut root = TreeNode::new(array[0]);
    for &value in &array[1..] {
        root.insert(value);
    }
    root
}

// 递归访问所有节点
fn tree_height(root: &Option<Box<TreeNode>>) -> usize {
    if let Some(root) = root.as_ref() {
        let left_height = tree_height(&root.left);
        let middle_height = tree_height(&root.middle);
        let right_height = tree_height(&root.right);
        let child_height = left_height.max(middle_height.max(right_height));
        1 + child_height
    } else {
        0
    }
}

fn main() {
    // 读取输入
    // 读取节点个数
    let mut line = String::new();
    let ret = stdin().lock().read_line(&mut line);
    assert!(ret.is_ok());
    let num_nodes: usize = line.trim().parse().unwrap();
    assert!((1..=10000).contains(&num_nodes));

    // 读取所有节点的值, 存放到数组
    line.clear();
    let ret = stdin().lock().read_line(&mut line);
    assert!(ret.is_ok());
    let array: Vec<i32> = line
        .trim()
        .split_ascii_whitespace()
        .map(|s| s.parse().unwrap())
        .collect();
    assert_eq!(array.len(), num_nodes);

    // 构造三叉树
    let root = build_tree(&array);
    let root = Some(Box::new(root));

    // 获取树的最大高度
    let height: usize = tree_height(&root);

    // 打印结果
    println!("{height}");
}