link
Method 1: use HashMap, Build undirected graph using treenodes as vertices, and BFS1
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// import java.util.*;
// Definition for a binary tree node.
/*
class TreeNode {
int val;
TreeNode left;
TreeNode right;
TreeNode(int x) {val = x;}
}*/
// Method 1: use HashMap
// 1. build a undirected graph using treenodes as vertices, and the parent-child relation as edges
// 2. do BFS with source vertice (target) to find all vertices with distance K to it.
class Solution {
Map<TreeNode, List<TreeNode>> map = new HashMap<>();
public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
List<Integer> res = new ArrayList<Integer>();
if (root == null || K < 0) return res;
buildGraph(root, null);
if (!map.containsKey(target)) return res;
Set<TreeNode> visited = new HashSet<TreeNode>();
Queue<TreeNode> q = new LinkedList<TreeNode>();
q.add(target);
visited.add(target);
while (!q.isEmpty()) {
int size = q.size();
if (K == 0) {
for (int i = 0; i < size; i++) res.add(q.poll().val);
return res;
}
for (int i = 0; i < size; i++) {
TreeNode node = q.poll();
for (TreeNode next : map.get(node)) {
if (visited.contains(next)) continue;
q.add(next);
visited.add(next);
}
}
K--;
}
return res;
}
private void buildGraph(TreeNode node, TreeNode parent) {
if (node == null) return;
if (!map.containsKey(node)) {
map.put(node, new ArrayList<TreeNode>());
if (parent != null) {
map.get(node).add(parent);
map.get(parent).add(node);
}
buildGraph(node.left, node);
buildGraph(node.right, node);
}
}
}
/*
public class MyClass {
public static void main(String[] args) {
// root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
TreeNode root = new TreeNode(3);
// LEFT BRANCH
root.left = new TreeNode(5);
root.left.left = new TreeNode(6);
root.left.right = new TreeNode(2);
root.left.right.left = new TreeNode(7);
root.left.right.right = new TreeNode(4);
// RIGHT BRANCH
root.right = new TreeNode(1);
root.right.left = new TreeNode(0);
root.right.right = new TreeNode(8);
TreeNode target = root.left;
int K = 2;
System.out.println(new Solution().distanceK(root, target, K));
System.out.println();
}
}*/
link
Method 2: No HashMap1
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//kind of like clone the tree, in the meanwhile add a parent link to the node
class Solution {
private GNode targetGNode;
private class GNode {
TreeNode node;
GNode parent, left, right;
GNode (TreeNode node) {
this.node = node;
}
}
public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
List<Integer> res = new ArrayList<Integer> ();
if (root == null || K < 0) return res;
cloneGraph(root, null, target);
if (targetGNode == null) return res;
Set<GNode> visited = new HashSet<GNode>();
Queue<GNode> q = new LinkedList<GNode>();
q.add(targetGNode);
visited.add(targetGNode);
while (!q.isEmpty()) {
int size = q.size();
if (K == 0) {
for (int i = 0; i < size ; i++) res.add(q.poll().node.val);
return res;
}
for (int i = 0; i < size; i++) {
GNode gNode = q.poll();
if (gNode.left != null && !visited.contains(gNode.left)) { visited.add(gNode.left); q.add(gNode.left); }
if (gNode.right != null && !visited.contains(gNode.right)) { visited.add(gNode.right); q.add(gNode.right); }
if (gNode.parent != null && !visited.contains(gNode.parent)) { visited.add(gNode.parent); q.add(gNode.parent); }
}
K--;
}
return res;
}
private GNode cloneGraph(TreeNode node, GNode parent, TreeNode target) {
if (node == null) return null;
GNode gNode = new GNode(node);
if (node == target) targetGNode = gNode;
gNode.parent = parent;
gNode.left = cloneGraph(node.left, gNode, target);
gNode.right = cloneGraph(node.right, gNode, target);
return gNode;
}
}