Another problem that can be solved easily using recursion and with the help of Kotlin’s operator, everything becomes easy-peasy.
/**
* 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
* }
*/
class Solution {
fun averageOfLevels(root: TreeNode?): DoubleArray {
var collector = HashMap<Int, ArrayList<TreeNode>>()
tailrec fun averageOfLevels_(node: TreeNode?, level: Int, collector: HashMap<Int, ArrayList<TreeNode>>){
if(node == null) return
if(!collector.containsKey(level)){
collector[level] = ArrayList<TreeNode>()
}
collector[level]?.add(node)
averageOfLevels_(node.left, level + 1, collector)
averageOfLevels_(node.right, level + 1, collector)
}
averageOfLevels_(root, 0, collector)
return collector.map { it.value?.map { it.`val`}.average() }.toDoubleArray()
}
}
Runtime: 244 ms, faster than 86.67% of Kotlin online submissions for Average of Levels in Binary Tree.
Memory Usage: 39.1 MB, less than 100.00% of Kotlin online submissions for Average of Levels in Binary Tree.
The main idea is to assign levels to each node. Then collect these levels into buckets (map) and finally loop through each bucket and get the average of the nodes in each bucket.