2265. Count Nodes Equal to Average of Subtree

count up if the node value is the same as the average of its own sub tree including itself.
and return the count

Kotlin

import kotlin.math.floor


data class Res(var sum: Int, var cnt: Int, var eq_cnt: Int)


class Solution {
    fun averageOfSubtree(root: TreeNode?): Int {
        fun dfs(node: TreeNode): Res {
            if (node.left == null && node.right == null) {
                return Res(node.`val`, 1, 1)
            }

            var sum: Int = 0
            var cnt: Int = 0
            var eq_cnt: Int = 0

            if (node.left != null) {
                var (ls, lc, le) = dfs(node.left)
                sum += ls
                cnt += lc
                eq_cnt += le
            }

            if (node.right != null) {
                var (rs, rc, re) = dfs(node.right)
                sum += rs
                cnt += rc
                eq_cnt += re
            }

            sum += node.`val`
            cnt += 1
            val avg = floor(sum.toDouble() / cnt.toDouble()).toInt()

            if (avg == node.`val`) {
                eq_cnt += 1
            }

            return Res(sum, cnt, eq_cnt)
        }

        if (root != null) {
            var (ls, lc, le) = dfs(root)
            return le
        }

        return 0
    }
}

Python

class Solution:
    def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
        def dfs(node):
            if not node.left and not node.right:
                return node.val, 1, 1

            ls, lc = 0, 0
            rs, rc = 0, 0
            le, re = 0, 0

            if node.left:
                ls, lc, le = dfs(node.left)

            if node.right:
                rs, rc, re = dfs(node.right)

            avg = (ls + rs + node.val) // (lc + rc + 1)

            eq_cnt = le + re
            if avg == node.val:
                eq_cnt += 1

            return ls + rs + node.val, lc + rc + 1, eq_cnt

        avg, cnt, eq_cnt = dfs(root)

        return eq_cnt