Family Tree

 1 class TreeNode {
 2     public String val;
 3     public TreeNode parent;
 4     public List<TreeNode> children;
 5 
 6     public TreeNode(String val) {
 7         this.val = val;
 8         children = new ArrayList<TreeNode>();
 9     }
10 }

 1 import java.util.*;
 2 
 3 public class Solution {
 4     public static void main(String[] args) {
 5         String s = "Frank->Mary,Mary->Sam,Mary->Bob,Sam->Katie,Sam->Pete";
 6         int target = 2;
 7         // Construct family tree
 8 
 9         TreeNode ancestor = new TreeNode("Ancestor");
10         Map<String, TreeNode> map = new HashMap<String, TreeNode>();
11         map.put("Ancestor", ancestor);
12 
13         String[] relations = s.split(",");
14 
15         for (String relation : relations) {
16             String[] names = relation.split("->");
17             String parent = names[0];
18             String child = names[1];
19             TreeNode parentNode, childNode;
20             if (map.containsKey(parent)) {
21                 parentNode = map.get(parent);
22             } else {
23                 parentNode = new TreeNode(parent);
24                 parentNode.parent = ancestor;
25             }
26             if (map.containsKey(child)) {
27                 childNode = map.get(child);
28             } else {
29                 childNode = new TreeNode(child);
30                 childNode.parent = parentNode;
31             }
32             List<TreeNode> childrenList = parentNode.children;
33             if (!childrenList.contains(childNode))
34                 childrenList.add(childNode);
35             map.put(parent, parentNode);
36             map.put(child, childNode);
37             System.out.println(parent);
38             System.out.println(child);
39 
40         }
41         // Find children of ancestor
42         List<TreeNode> firstChildren = ancestor.children;
43         for (String tmp : map.keySet()) {
44             TreeNode tmpNode = map.get(tmp);
45             if (tmpNode.parent == ancestor) {
46                 firstChildren.add(tmpNode);
47                 System.out.println(tmpNode.val);
48             }
49         }
50 
51         System.out.println("start BFS");
52         // BFS to get result
53         int level = 0;
54         List<TreeNode> currentList = new ArrayList<TreeNode>();
55         List<String> result = new ArrayList<String>();
56         List<TreeNode> nextList;
57         currentList.add(ancestor);
58         while (currentList.size() > 0) {
59             nextList = new ArrayList<TreeNode>();
60             for (TreeNode tmpNode : currentList) {
61                 List<TreeNode> childrenList = tmpNode.children;
62                 for (TreeNode oneChild : childrenList) {
63                     if (!nextList.contains(oneChild))
64                         nextList.add(oneChild);
65                 }
66                 currentList = nextList;
67                 level++;
68                 if (level == target) {
69                     for (TreeNode tmpNode2 : currentList)
70                         result.add(tmpNode2.val);
71                     break;
72                 }
73             }
74         }
75         for (String output : result) {
76             System.out.println(output);
77         }
78 
79     }
80 }