LeetCode Entry

2265. Count Nodes Equal to Average of Subtree

02.11.2023 medium 2023 kotlin

Number of nodes in a tree where val == sum / count of a subtree

2265. Count Nodes Equal to Average of Subtree medium blog post substack image.png

Join me on Telegram

https://t.me/leetcode_daily_unstoppable/390

Problem TLDR

Number of nodes in a tree where val == sum / count of a subtree

Intuition

Just do a Depth First Search and return sum and count of a subtree.

Approach

  • avoid nulls when traversing the tree

Complexity

  • Time complexity: \(O(n)\)

  • Space complexity: \(O(log(n))\) for the recursion depth

Code


    fun averageOfSubtree(root: TreeNode?): Int {
      var res = 0
      fun dfs(n: TreeNode): Pair<Int, Int> {
        val (ls, lc) = n.left?.let { dfs(it) } ?: 0 to 0
        val (rs, rc) = n.right?.let { dfs(it) } ?: 0 to 0
        val sum = n.`val` + ls + rs
        val count = 1 + lc + rc
        if (n.`val` == sum / count) res++
        return sum to count
      }
      root?.let { dfs(it) }
      return res
    }