算法题/树的BFS：连接层序兄弟结点

题目

给定一棵二叉树，将每个节点与其层次顺序的后继节点连接起来。每个级别的最后一个节点应该指向一个空节点。

解答

Java

import java.util.*;

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

  TreeNode(int x) {
    val = x;
    left = right = next = null;
  }

  //打印节点 level order traversal using 'next' pointer
  public void printLevelOrder() {
    TreeNode nextLevelRoot = this;
    while (nextLevelRoot != null) {
      TreeNode current = nextLevelRoot;
      nextLevelRoot = null;
      while (current != null) {
        System.out.print(current.val + " ");
        if (nextLevelRoot == null) {
          if (current.left != null)
            nextLevelRoot = current.left;
          else if (current.right != null)
            nextLevelRoot = current.right;
        }
        current = current.next;
      }
      System.out.println();
    }
  }
};

class ConnectLevelOrderSiblings {
  public static void connect(TreeNode root) {
    if (root == null)
      return;

    Queue<TreeNode> queue = new LinkedList<>();
    queue.offer(root);
    while (!queue.isEmpty()) {
      TreeNode previousNode = null;
      int levelSize = queue.size();
      //遍历这一层的所有节点
      for (int i = 0; i < levelSize; i++) {
        TreeNode currentNode = queue.poll();
        if (previousNode != null)
          previousNode.next = currentNode;
        previousNode = currentNode;

        //插入当前节点的子节点
        if (currentNode.left != null)
          queue.offer(currentNode.left);
        if (currentNode.right != null)
          queue.offer(currentNode.right);
      }
    }
  }

  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);
    ConnectLevelOrderSiblings.connect(root);
    System.out.println("Level order traversal using 'next' pointer: ");
    root.printLevelOrder();
  }
}

JavaScript

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

    //打印节点 level order traversal using 'next' pointer
    printLevelOrder() {
        let nextLevelRoot = this;
        while (nextLevelRoot !== null) {
            let current = nextLevelRoot;
            nextLevelRoot = null;
            while (current !== null) {
                console.log(current.val + " ");
                if (nextLevelRoot === null) {
                    if (current.left !== null)
                        nextLevelRoot = current.left;
                    else if (current.right !== null)
                        nextLevelRoot = current.right;
                }
                current = current.next;
            }
            console.log();
        }
    }
}

class ConnectLevelOrderSiblings {
    static connect(root) {
        if (root === null)
            return;

        let queue = [];
        queue.push(root);
        while (queue.length > 0) {
            let previousNode = null;
            let levelSize = queue.length;
            //遍历这一层的所有节点
            for (let i = 0; i < levelSize; i++) {
                let currentNode = queue.shift();
                if (previousNode !== null)
                    previousNode.next = currentNode;
                previousNode = currentNode;

                //插入当前节点的子节点
                if (currentNode.left !== null)
                    queue.push(currentNode.left);
                if (currentNode.right !== null)
                    queue.push(currentNode.right);
            }
        }
    }

    static main() {
        let 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);
        ConnectLevelOrderSiblings.connect(root);
        console.log("Level order traversal using 'next' pointer: ");
        root.printLevelOrder();
    }
}

ConnectLevelOrderSiblings.main();

Python

from collections import deque

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.next = None
  
    # 打印节点 level order traversal using 'next' pointer
    def printLevelOrder(self):
        nextLevelRoot = self
        while nextLevelRoot:
            current = nextLevelRoot
            nextLevelRoot = None
            while current:
                print(current.val, end=" ")
                if not nextLevelRoot:
                    if current.left:
                        nextLevelRoot = current.left
                    elif current.right:
                        nextLevelRoot = current.right
                current = current.next
            print()

class ConnectLevelOrderSiblings:
    def connect(root):
        if not root:
            return

        queue = deque()
        queue.append(root)
        while queue:
            previousNode = None
            levelSize = len(queue)
            # 遍历这一层的所有节点
            for _ in range(levelSize):
                currentNode = queue.popleft()
                if previousNode:
                    previousNode.next = currentNode
                previousNode = currentNode

                # 插入当前节点的子节点
                if currentNode.left:
                    queue.append(currentNode.left)
                if currentNode.right:
                    queue.append(currentNode.right)

    def main():
        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)
        ConnectLevelOrderSiblings.connect(root)
        print("Level order traversal using 'next' pointer: ")
        root.printLevelOrder()

ConnectLevelOrderSiblings.main()

Go

package main

import (
	"fmt"
)

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

//打印节点 level order traversal using 'next' pointer
func (node *TreeNode) printLevelOrder() {
	nextLevelRoot := node
	for nextLevelRoot != nil {
		current := nextLevelRoot
		nextLevelRoot = nil
		for current != nil {
			fmt.Printf("%d ", current.val)
			if nextLevelRoot == nil {
				if current.left != nil {
					nextLevelRoot = current.left
				} else if current.right != nil {
					nextLevelRoot = current.right
				}
			}
			current = current.next
		}
		fmt.Println()
	}
}

type ConnectLevelOrderSiblings struct{}

func (c *ConnectLevelOrderSiblings) connect(root *TreeNode) {
	if root == nil {
		return
	}

	queue := []*TreeNode{root}
	for len(queue) > 0 {
		previousNode := (*TreeNode)(nil)
		levelSize := len(queue)
		//遍历这一层的所有节点
		for i := 0; i < levelSize; i++ {
			currentNode := queue[0]
			queue = queue[1:]
			if previousNode != nil {
				previousNode.next = currentNode
			}
			previousNode = currentNode

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

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}

	connectSiblings := &ConnectLevelOrderSiblings{}
	connectSiblings.connect(root)

	fmt.Println("Level order traversal using 'next' pointer:")
	root.printLevelOrder()
}