2265. Count Nodes Equal to Average of Subtree¶

Intuition¶

Testing one node means knowing two numbers about its subtree: the sum of the values and the count of the nodes. Computing those separately for every node would rescan overlapping subtrees, but they compose upward for free — a node's totals are just its children's totals plus itself:

$$\text{sum}(v) = \text{sum}(v_{left}) + \text{sum}(v_{right}) + val(v)$$

$$\text{count}(v) = \text{count}(v_{left}) + \text{count}(v_{right}) + 1$$

So a single post-order traversal answers every node at once. Each call returns the pair to its parent and checks its own node on the way back up, which is why the whole tree is visited exactly once.

Approach: Post-Order DFS Returning `(sum, count)`¶

Define dfs(node) returning the pair for the subtree rooted at node:

An empty subtree returns (0, 0) — the identity for both accumulators, so leaves need no special case.
Recurse left and right to get their pairs.
Combine into this node's sum and count using the two identities above.
If val == sum / count (integer division), increment a counter that lives outside the recursion.
Return (sum, count) to the parent.

The order matters: the check happens after both recursive calls, because a node's own totals are not known until its children have reported. That is what makes this post-order rather than pre-order.

Why the counter is passed by reference¶

Each call already returns (sum, count), so the tally has to travel some other way. All three versions keep it outside the return value:

Python stores it on the instance (self.ans), which the nested dfs closes over.
Rust threads ans: &mut i32 through every call — the borrow checker allows this because only one call holds the mutable reference at a time.
Go passes result *int and does *result++.

That last line is worth a note: *result++ is valid Go and means (*result)++. Go's ++ is a statement applying to the whole expression, so there is none of C's *p++ ambiguity between incrementing the pointer and the pointee. Verified to compile and behave as intended.

Integer division is safe here, but only because values are non-negative¶

The problem asks for the average rounded down, and the three languages spell that differently: Python's // floors, while Rust and Go's / truncates toward zero. Those disagree on negatives — Python gives -7 // 2 == -4 where Go and Rust give -3.

It does not matter here: $$0 \le val \le 1000$$ means every subtree sum is non-negative and every count is positive, so flooring and truncation coincide. If the constraint ever admitted negative values the Python version would be the correct one and the other two would need div_euclid or an explicit adjustment.

No overflow¶

With at most 1000 nodes each at most 1000, the largest possible subtree sum is $$1000 \times 1000 = 10^6$$ — far inside a 32-bit integer, so Rust's i32 needs no widening. Confirmed by running the Rust version as a debug build, where an arithmetic overflow aborts.

The recursion depth is exactly at CPython's limit¶

The constraints allow a degenerate tree: 1000 nodes in a single spine, giving a recursion depth of 1000. CPython's default recursion limit is also 1000, and the Python version raises RecursionError: maximum recursion depth exceeded on that input when run with stock settings. Raising the limit to 10000 makes it return correctly.

LeetCode's judge raises the limit itself, so this passes there — but it is a real edge of the constraint space rather than a theoretical one. Go grows its goroutine stacks dynamically and Rust's 8 MB main-thread stack absorbs 1000 frames comfortably; neither is affected.

Rust: the borrow held across recursion¶

let n = node.borrow(); keeps a Ref alive while Self::dfs(&n.left, ..) and Self::dfs(&n.right, ..) run. That is fine because each child is a different RefCell, so no cell is borrowed twice at once. It would panic at runtime on a cyclic structure, which a binary tree never is. The debug build ran the whole test corpus without a borrow panic.

Worked example¶

root = [4,8,5,0,1,null,6]:

        4
      /   \
     8     5
    / \      \
   0   1      6

Post-order means children report before their parent:

visit order	node	subtree sum	count	`sum / count`	counts?
1	`0`	`0`	`1`	`0`	yes
2	`1`	`1`	`1`	`1`	yes
3	`8`	`9`	`3`	`9 / 3 = 3`	no
4	`6`	`6`	`1`	`6`	yes
5	`5`	`11`	`2`	`11 / 2 = 5`	yes
6	`4`	`24`	`6`	`24 / 6 = 4`	yes

Five nodes qualify. Two things this example demonstrates: every leaf always counts (a single value equals its own average), and node 5 only qualifies because the division floors — the exact average is 5.5.

For root = [1] the single node returns (1, 1) and 1 / 1 == 1, so the answer is 1.

Complexity¶

Time complexity: $$O(n)$$, where n is the number of nodes — each is visited once and does constant work.
Space complexity: $$O(h)$$ for the recursion stack, where h is the tree height. That is $$O(\log n)$$ when balanced and $$O(n)$$ for a degenerate spine, which the constraints permit.

Recomputing each subtree independently would be $$O(n \cdot h)$$ — up to $$O(n^2)$$ on a spine. The returned pair is what collapses that to linear.

Code¶

Go¶

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func averageOfSubtree(root *TreeNode) int {
    var dfs func (*TreeNode, *int) (int, int)
    dfs = func(root *TreeNode, result *int) (int, int) {
        if root == nil {
            return 0, 0
        }
        leftSum, leftCount := dfs(root.Left, result)
        rightSum, rightCount := dfs(root.Right, result)
        sum,count := leftSum + rightSum + root.Val, leftCount + rightCount + 1
        if root.Val == sum/count {
            *result++
        }
        return sum, count
    }
    ans := 0
    dfs(root, &ans)
    return ans
}

The var dfs func(...) declaration before the assignment is required: a closure cannot refer to itself inside a := initialiser, because the name is not in scope until that statement completes.

Rust¶

// Definition for a binary tree node.
// #[derive(Debug, PartialEq, Eq)]
// pub struct TreeNode {
//   pub val: i32,
//   pub left: Option<Rc<RefCell<TreeNode>>>,
//   pub right: Option<Rc<RefCell<TreeNode>>>,
// }
// 
// impl TreeNode {
//   #[inline]
//   pub fn new(val: i32) -> Self {
//     TreeNode {
//       val,
//       left: None,
//       right: None
//     }
//   }
// }
use std::rc::Rc;
use std::cell::RefCell;
impl Solution {
    pub fn dfs(root: &Option<Rc<RefCell<TreeNode>>>, ans: &mut i32) -> (i32, i32) {
        match root {
            Some(node) => {
                let n = node.borrow();
                let (left_sum, left_count) = Self::dfs(&n.left, ans);
                let (right_sum, right_count) = Self::dfs(&n.right, ans);
                let (sum, count) = (left_sum + right_sum + n.val, left_count + right_count + 1);
                if n.val == (sum / count) {
                    *ans += 1;
                }
                (sum, count)
            },
            None => (0, 0)
        }
    }

    pub fn average_of_subtree(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
        let mut ans = 0;
        Self::dfs(&root, &mut ans);
        ans
    }
}

Taking &Option<Rc<RefCell<TreeNode>>> rather than the owned value means the recursion borrows the tree instead of cloning Rc handles at every level, so the reference counts are never touched during the traversal.

Python¶

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def averageOfSubtree(self, root: TreeNode) -> int:
        self.ans = 0

        def dfs(root: TreeNode) -> (int, int):
            if root is None:
                return 0, 0
            left_sum, left_count = dfs(root.left)
            right_sum, right_count = dfs(root.right)
            _sum, count = left_sum + right_sum + root.val, left_count + right_count + 1
            if root.val == _sum // count:
                self.ans += 1

            return _sum, count

        dfs(root)
        return self.ans

_sum is spelled with a leading underscore to avoid shadowing the builtin sum, which is a habit worth keeping in a function that might later want to call it.

Test cases¶

tree	answer	what it exercises
`[4,8,5,0,1,null,6]`	`5`	Example 1 — traced above
`[1]`	`1`	Example 2 — single node
`[0]`	`1`	zero value, `0 / 1 == 0`
`[2,1,4]`	`3`	root `2` matches `7 / 3 = 2` by flooring
1000-node left spine of `7`s	`1000`	maximum depth; every node averages `7`
complete tree of 1000 nodes	—	maximum size

All three implementations were checked against a brute force that, for every node, walks that node's entire subtree from scratch and compares against the linear version. The corpus was 4005 trees: the two examples, 4000 randomly shaped trees of up to 40 nodes over value ranges 0..3, 0..10 and 0..1000 (small ranges make ties and near-misses common), a 1000-node complete tree, and a 1000-deep left spine. All three agreed on every tree, and the Rust build was made in debug mode where an overflow or a double RefCell borrow would panic — neither occurred.