aboutsummaryrefslogtreecommitdiff
path: root/kth_largest_sum_binary_tree.c
blob: f9597831573996347efb2b58a87c51a99dfef664 (plain)
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
// 2583. Kth Largest Sum in a Binary Tree
#include "leetcode.h"

/*
 * you are given the 'root' of a binary tree and a positive integer 'k'. the
 * level sum in the tree is the sum of the values of the nodes that are on the
 * same level. return the k'th largest level sum in the tree (not necessarily
 * distinct). if there are fewer than 'k' levels in the tree, return -1. note
 * that two nodes are on the same level if they have the same distance from the
 * root.
 */

int depth(struct TreeNode *root) {
  if (!root)
    return 0;
  int l = depth(root->left) + 1;
  int r = depth(root->right) + 1;
  return l > r ? l : r;
}

void level_sum(struct TreeNode *root, int level, long long *data) {
  if (!root)
    return;
  data[level] += root->val;
  level_sum(root->left, level + 1, data);
  level_sum(root->right, level + 1, data);
}

void swap(long long *a, long long *b) {
  long long tmp = *a;
  *a = *b;
  *b = tmp;
}

int partition(long long *data, int l, int h) {
  int i = l, j = l;
  long long pivot = data[h - 1];
  for (j - l; j < h - 1; j++)
    if (data[j] < pivot)
      swap(&data[i++], &data[j]);
  swap(&data[i], &data[h - 1]);
  return i;
}

void quick_select(long long *data, int l, int h, int k) {
  if (l + 1 < h) {
    int p = partition(data, l, h);
    if (p == k)
      return;
    else if (p < k)
      quick_select(data, p + 1, h, k);
    else
      quick_select(data, l, p, k);
  }
}

long long kthLargestLevelSum(struct TreeNode *root, int k) {
  int n = depth(root);
  if (k > n)
    return -1;
  long long val, *data = (long long *)calloc(n, sizeof(long long));
  level_sum(root, 0, data);
  quick_select(data, 0, n, n - k);
  return data[n - k];
}

int main() {
  struct TreeNode *r1 = treenode_create(5);
  r1->left = treenode_create(8);
  r1->right = treenode_create(9);
  r1->left->left = treenode_create(2);
  r1->left->right = treenode_create(1);
  r1->right->left = treenode_create(3);
  r1->right->right = treenode_create(7);
  r1->left->left->left = treenode_create(4);
  r1->left->left->right = treenode_create(6);
  struct TreeNode *r2 = treenode_create(1);
  r2->left = treenode_create(2);
  r2->right = NULL;
  r2->left->left = treenode_create(3);
  printf("%lld\n", kthLargestLevelSum(r1, 2)); // expect: 13
  printf("%lld\n", kthLargestLevelSum(r2, 1)); // expect: 3
  treenode_free(r1);
  treenode_free(r2);
}