跳至主要內容

116, 填充每个节点的下一个右侧节点指针

Mike大约 6 分钟binary treemediumbinary treelinked listdepth first searchbreadth first search

一、题目描述

给定一个 完美二叉树 ，其所有叶子节点都在同一层，每个父节点都有两个子节点。二叉树定义如下：

struct Node {
    int val;
    Node *left;
    Node *right;
    Node *next;
}

填充它的每个next指针，让这个指针指向其下一个右侧节点。如果找不到下一个右侧节点，则将next指针设置为NULL。

初始状态下，所有next指针都被设置为NULL。

示例 1
picture116
输入: root = [1, 2, 3, 4, 5, 6, 7]
输出: [1, #, 2, 3, #, 4, 5, 6, 7, #]
解释: 给定二叉树如图A所示，你的函数应该填充它的每个next指针，以指向其下一个右侧节点，如图B所示。序列化的输出按层序遍历排列，同一层节点由next指针连接，'#'标志着每一层的结束。

示例 2
输入: root = []
输出: []

提示

  • 树中节点的数量在[0, 2¹² - 1]范围内
  • -1000 <= node.val <= 1000

进阶

  • 你只能使用常量级额外空间。
  • 使用递归解题也符合要求，本题中递归程序占用的栈空间不算做额外的空间复杂度。

相关主题

  • 树
  • 深度优先搜索
  • 广度优先搜索
  • 链表
  • 二叉树

二、题解

#[derive(Debug, PartialEq, Eq)]
pub struct Node {
    pub val: i32,
    pub left: Option<Rc<RefCell<Node>>>,
    pub right: Option<Rc<RefCell<Node>>>,
    pub next: Option<Rc<RefCell<Node>>>,
}

impl Node {
    pub fn new(val: i32) -> Option<Rc<RefCell<Node>>> {
        Some(Rc::new(RefCell::new(Node {
            val,
            left: None,
            right: None,
            next: None,
        })))
    }
    pub fn new_with_children(
        val: i32,
        left: Option<Rc<RefCell<Node>>>,
        right: Option<Rc<RefCell<Node>>>,
    ) -> Option<Rc<RefCell<Node>>> {
        Some(Rc::new(RefCell::new(Node {
            val,
            left,
            right,
            next: None,
        })))
    }
}

方法 1: 层序遍历

pub fn connect(root: Option<Rc<RefCell<Node>>>) -> Option<Rc<RefCell<Node>>> {
    //Self::bfs_iter_1(root)
    Self::bfs_iter_2(root)
}

fn bfs_iter_1(root: Option<Rc<RefCell<Node>>>) -> Option<Rc<RefCell<Node>>> {
    if let Some(root) = root.clone() {
        let mut queue = VecDeque::from([root]);

        while !queue.is_empty() {
            let level_len = queue.len();
            let mut prev: Option<Rc<RefCell<Node>>> = None;

            for i in 0..level_len {
                if let Some(curr) = queue.pop_front() {
                    if i > 0 {
                        prev.map(|prev| {
                            prev.borrow_mut().next = Some(curr.clone());
                        });
                    }
                    prev = Some(curr.clone());
                    if let Some(left) = curr.borrow().left.clone() {
                        queue.push_back(left);
                    }
                    if let Some(right) = curr.borrow().right.clone() {
                        queue.push_back(right);
                    }
                };
            }
        }
    }

    root
}

fn bfs_iter_2(root: Option<Rc<RefCell<Node>>>) -> Option<Rc<RefCell<Node>>> {
    if let Some(root) = root.clone() {
        let mut queue = VecDeque::from([(root.clone(), 0)]);
        let mut prev: (Option<Rc<RefCell<Node>>>, i32) = (None, -1);

        while let Some((curr, level)) = queue.pop_front() {
            if prev.1 == level {
                prev.0.map(|prev| {
                    prev.borrow_mut().next = Some(curr.clone());
                });
            };
            prev = (Some(curr.clone()), level);

            if let Some(left) = curr.borrow().left.clone() {
                queue.push_back((left, level + 1));
            }
            if let Some(right) = curr.borrow().right.clone() {
                queue.push_back((right, level + 1));
            }
        }
    }

    root
}

方法 2: 使用已建立的next指针

pub fn connect(root: Option<Rc<RefCell<Node>>>) -> Option<Rc<RefCell<Node>>> {
    //Self::use_next_pointer_iter(root)
    Self::use_next_pointer_recur(root)
}
///
/// 迭代版
///
fn use_next_pointer_iter(root: Option<Rc<RefCell<Node>>>) -> Option<Rc<RefCell<Node>>> {
    let mut leftmost = root.clone();

    while let Some(level_first) = leftmost {
        let mut level = Some(level_first.clone());

        while let Some(curr) = level {
            match (curr.borrow().left.clone(), curr.borrow().right.clone()) {
                (Some(left), Some(right)) => {
                    left.borrow_mut().next = Some(right.clone());
                    if let Some(next) = curr.borrow().next.clone() {
                        right.borrow_mut().next = next.borrow().left.clone();
                    }
                }
                (_, _) => break,
            }
            level = curr.borrow().next.clone();
        }

        leftmost = level_first.borrow().left.clone();
    }

    root
}
///
/// 递归版(前序)
///
fn use_next_pointer_recur(root: Option<Rc<RefCell<Node>>>) -> Option<Rc<RefCell<Node>>> {
    const PRE_ORDER: fn(Option<Rc<RefCell<Node>>>) = |root| {
        if let Some(curr) = root {
            match (curr.borrow().left.clone(), curr.borrow().right.clone()) {
                (Some(left), Some(right)) => {
                    left.borrow_mut().next = Some(right.clone());
                    if let Some(next) = curr.borrow().next.clone() {
                        right.borrow_mut().next = next.borrow().left.clone();
                    }
                }
                (_, _) => return,
            }

            PRE_ORDER(curr.borrow().left.clone());
            PRE_ORDER(curr.borrow().right.clone());
        }
    };

    PRE_ORDER(root.clone());

    root
}