rust
复制代码
// [3,9,20,null,null,15,7] true
// [1,2,2,3,3,null,null,4,4] false
// [] true
// Definition for a binary tree node.
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Option<Rc<RefCell<TreeNode>>>,
pub right: Option<Rc<RefCell<TreeNode>>>,
}
impl TreeNode {
#[inline]
pub fn new(val: i32) -> Self {
TreeNode {
val,
left: None,
right: None
}
}
}
use std::rc::Rc;
use std::cell::RefCell;
pub fn is_balanced(root: Option<Rc<RefCell<TreeNode>>>) -> bool {
// Helper function to determine height and if the subtree is balanced
fn check_height_and_balance(node: &Option<Rc<RefCell<TreeNode>>>) -> (bool, i32) {
if let Some(node) = node {
let left = check_height_and_balance(&node.borrow().left);
let right = check_height_and_balance(&node.borrow().right);
let is_balanced = left.0 && right.0 && (left.1 - right.1).abs() <= 1;
let height = 1 + left.1.max(right.1);
(is_balanced, height)
} else {
(true, 0)
}
}
check_height_and_balance(&root).0
}
fn build_tree(values: Vec<Option<i32>>) -> Option<Rc<RefCell<TreeNode>>> {
if values.is_empty() {
return None;
}
let root = Rc::new(RefCell::new(TreeNode::new(values[0].unwrap())));
let mut queue = vec![Rc::clone(&root)];
let mut i = 1;
while i < values.len() {
let current = queue.remove(0);
if let Some(val) = values[i] {
let left_node = Rc::new(RefCell::new(TreeNode::new(val)));
current.borrow_mut().left = Some(Rc::clone(&left_node));
queue.push(left_node);
}
i += 1;
if i < values.len() {
if let Some(val) = values[i] {
let right_node = Rc::new(RefCell::new(TreeNode::new(val)));
current.borrow_mut().right = Some(Rc::clone(&right_node));
queue.push(right_node);
}
i += 1;
}
}
Some(root)
}
fn main() {
let tree = build_tree(vec![Some(3), Some(9), Some(20), None, None, Some(15), Some(7)]);
assert_eq!(is_balanced(tree), true);
let tree = build_tree(vec![Some(1), Some(2), Some(2), Some(3), Some(3), None, None, Some(4), Some(4)]);
assert_eq!(is_balanced(tree), false);
let tree = build_tree(vec![]);
assert_eq!(is_balanced(tree), true);
}