Leetcode 1973 Count Nodes Equal to Sum of Descendants

Problem statement

https://leetcode.com/problems/count-nodes-equal-to-sum-of-descendants/

Solution

Just keep function dfs(node), which returns sum of all nodes in subtree and calculate it recursively.

Complexity

Time complexity is O(n), space complexity is O(h).

Code

class Solution:
    def equalToDescendants(self, root):
        self.ans = 0
        
        def dfs(node):
            if not node: return 0
            sm = dfs(node.left) + dfs(node.right)
            if sm == node.val: self.ans += 1
            return sm + node.val
        
        dfs(root)
        return self.ans