leetcode

Leetcode Solutions (0x6A73)
git clone https://git.pngu.org/leetcode
Log | Files | Refs | Submodules | README

cnt_node_equal_subtree.c (1047B)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
// 2265. Count Nodes Equal to Average of Subtree
#include "leetcode.h"

/*
 * given the root of a binary tree, return the number of nodes where the value
 * of the node is equal to the average of the values in its subtree. note, the
 * average of 'n' elements is the sum of the n elements divided by n and rounded
 * down to the nearest integer. also, a subtree of root is a tree consisting of
 * root and all of its descendants
 */

struct TreeNode {
  int val;
  struct TreeNode *left;
  struct TreeNode *right;
};

int cnt_subtree(struct TreeNode *root, int *res) {
  if (!root)
    return 0;
  int left = cnt_subtree(root->left, res);
  int rght = cnt_subtree(root->right, res);
  int curr_cnt = left + rght + 1, curr_sum = root->val;
  curr_sum += (root->left) ? root->left->val : 0;
  curr_sum += (root->right) ? root->right->val : 0;
  if (root->val == (curr_sum / curr_cnt))
    (*res)++;
  root->val = curr_sum;
  return curr_cnt;
}

int averageOfSubtree(struct TreeNode *root) {
  int ans = 0;
  cnt_subtree(root, &ans);
  return ans;
}