算法题/树的BFS：层序继承者

题目

给定一棵二叉树和一个节点，找出树中给定节点的级序后继子节点。级别顺序后继节点是在级别顺序遍历中出现在给定节点之后的节点。

解答

Java

import java.util.*;

class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;

  TreeNode(int x) {
    val = x;
  }
};

class LevelOrderSuccessor {
  public static TreeNode findSuccessor(TreeNode root, int key) {
    if (root == null)
      return null;

    Queue<TreeNode> queue = new LinkedList<>();
    queue.offer(root);
    while (!queue.isEmpty()) {
      TreeNode currentNode = queue.poll();
      //插入当前节点的子节点到队列中
      if (currentNode.left != null)
        queue.offer(currentNode.left);
      if (currentNode.right != null)
        queue.offer(currentNode.right);

      //如果找到key,则跳出循环
      if (currentNode.val == key)
        break;
    }

    return queue.peek();//跳出循环后，下一个节点就是要找的节点
  }

  public static void main(String[] args) {
    TreeNode root = new TreeNode(12);
    root.left = new TreeNode(7);
    root.right = new TreeNode(1);
    root.left.left = new TreeNode(9);
    root.right.left = new TreeNode(10);
    root.right.right = new TreeNode(5);
    TreeNode result = LevelOrderSuccessor.findSuccessor(root, 12);
    if (result != null)
      System.out.println(result.val + " ");
    result = LevelOrderSuccessor.findSuccessor(root, 9);
    if (result != null)
      System.out.println(result.val + " ");
  }
}

JavaScript

class TreeNode {
  constructor(x) {
    this.val = x;
    this.left = null;
    this.right = null;
  }
}

class LevelOrderSuccessor {
  static findSuccessor(root, key) {
    if (root === null)
      return null;

    const queue = [];
    queue.push(root);
    while (queue.length > 0) {
      const currentNode = queue.shift();
      // 将当前节点的子节点加入队列
      if (currentNode.left !== null)
        queue.push(currentNode.left);
      if (currentNode.right !== null)
        queue.push(currentNode.right);

      // 如果找到 key，则跳出循环
      if (currentNode.val === key)
        break;
    }

    return queue[0]; // 跳出循环后，下一个节点就是要找的节点
  }
}

const root = new TreeNode(12);
root.left = new TreeNode(7);
root.right = new TreeNode(1);
root.left.left = new TreeNode(9);
root.right.left = new TreeNode(10);
root.right.right = new TreeNode(5);

let result = LevelOrderSuccessor.findSuccessor(root, 12);
if (result !== null)
  console.log(result.val);
  
result = LevelOrderSuccessor.findSuccessor(root, 9);
if (result !== null)
  console.log(result.val);

Python

from collections import deque

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class LevelOrderSuccessor:
    def findSuccessor(root, key):
        if root is None:
            return None

        queue = deque()
        queue.append(root)
        while queue:
            currentNode = queue.popleft()
            # 将当前节点的子节点加入队列
            if currentNode.left:
                queue.append(currentNode.left)
            if currentNode.right:
                queue.append(currentNode.right)

            # 如果找到 key，则跳出循环
            if currentNode.val == key:
                break

        if queue:
            return queue[0]  # 跳出循环后，下一个节点就是要找的节点
        else:
            return None

root = TreeNode(12)
root.left = TreeNode(7)
root.right = TreeNode(1)
root.left.left = TreeNode(9)
root.right.left = TreeNode(10)
root.right.right = TreeNode(5)

result = LevelOrderSuccessor.findSuccessor(root, 12)
if result is not None:
    print(result.val)

result = LevelOrderSuccessor.findSuccessor(root, 9)
if result is not None:
    print(result.val)

Go

package main

import (
	"fmt"
)

type TreeNode struct {
	val   int
	left  *TreeNode
	right *TreeNode
}

func findSuccessor(root *TreeNode, key int) *TreeNode {
	if root == nil {
		return nil
	}

	queue := []*TreeNode{root}
	for len(queue) > 0 {
		currentNode := queue[0]
		queue = queue[1:]

		// 将当前节点的子节点加入队列
		if currentNode.left != nil {
			queue = append(queue, currentNode.left)
		}
		if currentNode.right != nil {
			queue = append(queue, currentNode.right)
		}

		// 如果找到 key，则跳出循环
		if currentNode.val == key {
			break
		}
	}

	if len(queue) > 0 {
		return queue[0] // 跳出循环后，下一个节点就是要找的节点
	} else {
		return nil
	}
}

func main() {
	root := &TreeNode{val: 12}
	root.left = &TreeNode{val: 7}
	root.right = &TreeNode{val: 1}
	root.left.left = &TreeNode{val: 9}
	root.right.left = &TreeNode{val: 10}
	root.right.right = &TreeNode{val: 5}

	result := findSuccessor(root, 12)
	if result != nil {
		fmt.Println(result.val)
	}

	result = findSuccessor(root, 9)
	if result != nil {
		fmt.Println(result.val)
	}
}