Следующий узел при обходе по уровням

Следующий узел при обходе по уровням


Условие задачи

Даны бинарное дерево и значение узла. Найдите следующий узел при обходе дерева по уровням — узел, который появляется сразу после заданного в последовательности BFS.

Следующий узел при обходе бинарного дерева по уровням

Решение

Задача следует паттерну обхода бинарного дерева по уровням. Нам не нужно сохранять уровни: достаточно извлекать узлы из очереди и сразу добавлять их потомков. Когда найден узел с заданным значением, следующий элемент в очереди и будет его преемником при BFS-обходе.

Если заданный узел идёт последним или не найден, функция вернёт nil.

Код

Вот как будет выглядеть наш алгоритм:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
package main

import "fmt"

type TreeNode struct {
	Value 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)
		}

		if currentNode.Value == key {
			break
		}
	}

	if len(queue) == 0 {
		return nil
	}
	return queue[0]
}

func printSuccessor(root *TreeNode, key int) {
	successor := findSuccessor(root, key)
	if successor == nil {
		fmt.Printf("После %d нет следующего узла\n", key)
		return
	}
	fmt.Printf("Следующий узел после %d: %d\n", key, successor.Value)
}

func main() {
	root := &TreeNode{Value: 12}
	root.Left = &TreeNode{Value: 7}
	root.Right = &TreeNode{Value: 1}
	root.Left.Left = &TreeNode{Value: 9}
	root.Right.Left = &TreeNode{Value: 10}
	root.Right.Right = &TreeNode{Value: 5}

	printSuccessor(root, 12)
	printSuccessor(root, 9)
}

Вывод:

1
2
Следующий узел после 12: 7
Следующий узел после 9: 10

Временная сложность

В худшем случае алгоритм посещает все узлы, поэтому его временная сложность равна \(O(N)\).

Пространственная сложность

Для очереди требуется \(O(W)\) памяти, где \(W\) — максимальная ширина дерева. В худшем случае пространственная сложность равна \(O(N)\).