Sponsored
Sponsored
This approach involves using a post-order traversal (left, right, root) to compute the sum and count of nodes for each subtree. At each node, the sum of values and the total count of the subtree nodes are determined. The average is calculated by dividing the sum by the count and compared with the node's value. If they are equal, a counter is incremented.
Time Complexity: O(n), where n is the number of nodes in the tree, as each node is visited once.
Space Complexity: O(h), where h is the height of the tree, which accounts for the recursion stack.
1class TreeNode:
2 def __init__(self, val=0, left=None, right=None):
3 self.val = val
4 self.left = left
5 self.right = right
6
7class Solution:
8 def averageOfSubtree(self, root: TreeNode) -> int:
9 def dfs(node):
10 if not node:
11 return (0, 0)
12 left_sum, left_count = dfs(node.left)
13 right_sum, right_count = dfs(node.right)
14 total_sum = left_sum + right_sum + node.val
15 total_count = left_count + right_count + 1
16 if node.val == total_sum // total_count:
17 self.count += 1
18 return total_sum, total_count
19 self.count = 0
20 dfs(root)
21 return self.count
This Python solution defines a recursive DFS function `dfs` that traverses the tree in a post-order manner. For each node, it computes the sum and count of the subtree rooted at that node. If the node value equals the average (sum // count) of its subtree, the counter `self.count` is incremented.
This approach involves using an iterative method with an explicit stack to perform a post-order traversal. Similar to the recursive method, it calculates the sum and count of the subtree nodes for each node using an additional data structure to simulate recursion.
Time Complexity: O(n), since each node is processed twice.
Space Complexity: O(n), for the stack and the dictionary storing sums and counts for each node.
1public class TreeNode {
2 public int val;
3 public TreeNode left;
4 public TreeNode right;
5 public TreeNode(int val = 0, TreeNode left = null, TreeNode right = null) {
this.val = val;
this.left = left;
this.right = right;
}
}
public class Solution {
public int AverageOfSubtree(TreeNode root) {
int count = 0;
Stack<(TreeNode, bool)> stack = new Stack<(TreeNode, bool)>();
Dictionary<TreeNode, (int, int)> nodeData = new Dictionary<TreeNode, (int, int)>();
stack.Push((root, false));
while (stack.Count > 0) {
var (node, visited) = stack.Pop();
if (node != null) {
if (visited) {
int leftSum = node.left != null ? nodeData[node.left].Item1 : 0;
int leftCount = node.left != null ? nodeData[node.left].Item2 : 0;
int rightSum = node.right != null ? nodeData[node.right].Item1 : 0;
int rightCount = node.right != null ? nodeData[node.right].Item2 : 0;
int totalSum = leftSum + rightSum + node.val;
int totalCount = leftCount + rightCount + 1;
nodeData[node] = (totalSum, totalCount);
if (node.val == totalSum / totalCount) count++;
} else {
stack.Push((node, true));
stack.Push((node.right, false));
stack.Push((node.left, false));
}
}
}
return count;
}
}
The C# solution utilizes an explicit stack to perform a post-order traversal. It uses a dictionary to store the sum and count for each node. The node is processed twice: first to traverse its children and then to calculate its subtree's sum and count.