aboutsummaryrefslogtreecommitdiff
path: root/all_nodes_binary_tree.c
blob: 05e1115c97d9e013c0ced8c8696e710f39692d5b (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
// 863. All Nodes Distance K in Binary Tree
#include <math.h>
#include <stdlib.h>

/*
 * given the 'root' of a binary tree, the value of a target node 'target', and
 * an integer 'k', return an array of the values of all nodes that have a
 * distance 'k' from the target node. you can return the answer in any order.
 */

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

int depth(struct TreeNode *root) {
  if (!root)
    return -1;
  return fmax(depth(root->left), depth(root->right)) + 1;
}

void find_path(struct TreeNode *root, struct TreeNode *target,
               struct TreeNode **data, int data_idx, struct TreeNode **path,
               int *path_idx) {
  if (!root)
    return;
  data[data_idx] = root;
  if (root == target) {
    *path_idx = data_idx + 1;
    for (int i = 0; i < data_idx + 1; i++)
      path[i] = data[i];
    return;
  }
  find_path(root->left, target, data, data_idx + 1, path, path_idx);
  find_path(root->right, target, data, data_idx + 1, path, path_idx);
}

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

int *distanceK(struct TreeNode *root, struct TreeNode *target, int k,
               int *return_size) {
  int *ans = malloc(100 * sizeof(int)), idx = 0;
  if (root == target) {
    k_level_data(root, ans, &idx, k);
    *return_size = idx;
    return ans;
  }
  int dep = depth(root);
  struct TreeNode **data = malloc((dep + 1) * sizeof(struct TreeNode *));
  struct TreeNode **path = malloc((dep + 1) * sizeof(struct TreeNode *));
  int data_idx = 0, path_idx = 0, a = 1;
  find_path(root, target, data, data_idx, path, &path_idx);
  k_level_data(target, ans, &idx, k);
  for (int i = path_idx - 2; i >= 0; i--) {
    if (!(k - a)) {
      ans[idx] = path[i]->val;
      idx++;
      break;
    }
    if (path[i]->left == path[i + 1])
      k_level_data(path[i]->right, ans, &idx, k - a - 1);
    else
      k_level_data(path[i]->left, ans, &idx, k - a - 1);
    a++;
  }
  *return_size = idx;
  return ans;
}