Count Complete Tree Nodes

Условие

Дано полное двоичное дерево, подсчитайте количество его узлов.

Примеры

Input: root = [1, 2, 3, 4, 5, 6]

Output: 6

Input: root = []

Output: 0

Input: root = [1]

Output: 1

Решение

/**
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
fun countNodes(root: TreeNode?): Int {
    // Если корень пустой, возвращаем 0
    if (root == null) return 0

    // Локальная функция для вычисления глубины поддерева
    fun getDepth(node: TreeNode?): Int {
        var depth = 0
        var current = node
        while (current != null) {
            depth++
            current = current.left // Идем по левой ветви
        }
        return depth
    }

    val leftDepth = getDepth(root.left) // Получаем глубину левого поддерева
    val rightDepth = getDepth(root.right) // Получаем глубину правого поддерева

    // Если глубины равны, значит, левое поддерево полно и мы можем использовать формулу
    return if (leftDepth == rightDepth) {
        (1 shl leftDepth) + countNodes(root.right) // 2^leftDepth + количество узлов в правом поддереве
    } else {
        (1 shl rightDepth) + countNodes(root.left) // 2^rightDepth + количество узлов в левом поддереве
    }
}

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

O(log^2 n), где n — количество узлов, так как мы вычисляем глубину каждого поддерева и используем это для вычисления количества узлов.

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

O(log n), в худшем случае, для стека рекурсии при вызовах функций.