Условие

Дано бинарное дерево. Найдите среднее значение элементов на каждом уровне дерева и верните его в виде массива.

Примеры

Input: root = [3, 9, 20, null, null, 15, 7] Output: [3.00000, 14.50000, 11.00000] Explanation: Среднее на уровне 0 = 3, на уровне 1 = (9 + 20) / 2 = 14.5, на уровне 2 = (15 + 7) / 2 = 11.

Input: root = [3, 9, 20, 15, 7] Output: [3.00000, 14.50000, 11.00000]

Решение

/**
 * 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 averageOfLevels(root: TreeNode?): DoubleArray {
    val result = mutableListOf<Double>()
    if (root == null) return doubleArrayOf()

    val queue: ArrayDeque<TreeNode> = ArrayDeque()
    queue.add(root)

    while (queue.isNotEmpty()) {
        val levelSize = queue.size
        var sum = 0.0

        for (i in 0 until levelSize) {
            val node = queue.removeFirst()
            sum += node.`val`
            node.left?.let { queue.add(it) } // Добавляем левое поддерево, если есть
            node.right?.let { queue.add(it) } // Добавляем правое поддерево, если есть
        }

        result.add(sum / levelSize) // Вычисляем среднее для текущего уровня
    }

    return result.toDoubleArray()
}

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

O(n), где n — количество узлов в дереве, так как каждый узел обрабатывается один раз.

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

O(m), где m — максимальное количество узлов на одном уровне дерева (ширина дерева).