Leetcode

Equal Tree Partition

  • Time:O(n)
  • Space:O(n)

C++

class Solution {
 public:
  bool checkEqualTree(TreeNode* root) {
    if (!root)
      return false;

    unordered_set<int> seen;
    const int sum = root->val + dfs(root->left, seen) + dfs(root->right, seen);
    return sum % 2 == 0 && seen.count(sum / 2);
  }

 private:
  int dfs(TreeNode* root, unordered_set<int>& seen) {
    if (!root)
      return 0;

    const int sum = root->val + dfs(root->left, seen) + dfs(root->right, seen);
    seen.insert(sum);
    return sum;
  }
};

JAVA

class Solution {
  public boolean checkEqualTree(TreeNode root) {
    if (root == null)
      return false;

    Set<Integer> seen = new HashSet<>();
    final int sum = root.val + dfs(root.left, seen) + dfs(root.right, seen);
    return sum % 2 == 0 && seen.contains(sum / 2);
  }

  private int dfs(TreeNode root, Set<Integer> seen) {
    if (root == null)
      return 0;

    final int sum = root.val + dfs(root.left, seen) + dfs(root.right, seen);
    seen.add(sum);
    return sum;
  }
}

Python

class Solution:
  def checkEqualTree(self, root: Optional[TreeNode]) -> bool:
    if not root:
      return False

    seen = set()

    def dfs(root: Optional[TreeNode]) -> int:
      if not root:
        return 0

      sum = root.val + dfs(root.left) + dfs(root.right)
      seen.add(sum)
      return sum

    sum = root.val + dfs(root.left) + dfs(root.right)
    return sum % 2 == 0 and sum // 2 in seen