Average of Levels in Binary Tree (LeetCode)
Here is the problem (link: https://leetcode.com/problems/average-of-levels-in-binary-tree/#/description):
Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.
Example 1:
Input: 3 / \ 9 20 / \ 15 7 Output: [3, 14.5, 11] Explanation: The average value of nodes on level 0 is 3, on level 1 is 14.5, and on level 2 is 11. Hence return [3, 14.5, 11].
Note:
- The range of node's value is in the range of 32-bit signed integer.
One way to solve it is the following:
a) First get the number of levels in the tree (one-line function)
b) After that do a DFS where you pass along an array to keep track of the sum and count per level, and an index indicating in which level you're at
c) At the end transform the array into an IList while calculating the avg per level
Cheers, Marcelo.
public class Solution
{
public IList<double> AverageOfLevels(TreeNode root)
{
if (root == null) return new List<double>();
AverageLevel[] avg = new AverageLevel[NumberOfLevels(root)];
for (int i = 0; i < avg.Length; i++) avg[i] = new AverageLevel();
AverageOfLevelsInternal(root, avg, 0);
List<double> solution = new List<double>();
for (int i = 0; i < avg.Length; i++)
{
solution.Add(avg[i].sum * 1.0 / avg[i].count);
}
return solution;
}
private void AverageOfLevelsInternal(TreeNode root, AverageLevel[] avg, int currentIndex)
{
if (root == null) return;
avg[currentIndex].sum += root.val;
avg[currentIndex].count++;
AverageOfLevelsInternal(root.left, avg, currentIndex + 1);
AverageOfLevelsInternal(root.right, avg, currentIndex + 1);
}
private class AverageLevel
{
public double sum;
public int count;
}
private int NumberOfLevels(TreeNode root)
{
return (root == null) ? 0 : (Math.Max(NumberOfLevels(root.left), NumberOfLevels(root.right)) + 1);
}
}
nice DFS recursive solution. I prefer BFS though, because this way it's easy to track sum and count using local variables.
ReplyDeleteMy take on it:
class Solution {
public:
vector averageOfLevels(const TreeNode* root) {
if (root == nullptr) return {};
queue current, next;
current.push(root);
vector averages;
long sum = 0; int count = 0;
while (!current.empty()) {
auto node = current.front(); current.pop();
sum += node->val; count += 1;
if (node->left != nullptr) next.push(node->left);
if (node->right != nullptr) next.push(node->right);
if (current.empty()) {
swap(current, next);
averages.push_back(1.0 * sum / count);
sum = 0; count = 0;
}
}
return averages;
}
};
Keep them coming Marcelo!
Nice!!!
ReplyDelete