leetcode

0 class Solution {

1 TreeNode *parent[501];

2 void find(TreeNode *root, TreeNode *target) {

3 queue<TreeNode *> q({root});

4

5 parent[root->val] = NULL;

6 while (!q.empty()) {

7 for (int k = q.size(); k > 0; k--) {

8 TreeNode *root = q.front();

9 q.pop();

10 if (root == target) return;

11 if (root->left) {

12 parent[root->left->val] = root;

13 q.push(root->left);

14 }

15 if (root->right) {

16 parent[root->right->val] = root;

17 q.push(root->right);

18 }

19 }

20 }

21 }

22

23 public:

24 vector<int> distanceK(TreeNode *root, TreeNode *target, int k) {

25 if (k == 0) return {target->val};

26 find(root, target);

27

28 vector<int> res;

29 while (k-- >= 0) {

30 queue<TreeNode *> q({target});

31 for (int i = 0; i <= k && !q.empty(); i++) {

32 for (int t = q.size(); t > 0; t--) {

33 TreeNode *root = q.front();

34 q.pop();

35 if (root->left) q.push(root->left);

36 if (root->right) q.push(root->right);

37 }

38 }

39 while (!q.empty()) {

40 res.push_back(q.front()->val);

41 q.pop();

42 }

43

44 TreeNode *p = parent[target->val];

45 if (!p) break;

46 (p->left == target ? p->left : p->right) = NULL;

47 target = p;

48 }

49 return res;

50 }

51 };