Family Tree

Question

A traditional constructing tree problem.

Given a string to represent relationships, and print out number n level names.

For example,

Input: "Frank->Mary,Mary->Sam,Mary->Bob,Sam->Katie,Sam->Pete", 2

Output: [Mary]

Solution

Programming thinking is simple, first, we need to construct a family tree according to input, then we do BFS to find results.

Two trick points to notice.

1. Create an ancestor treenode first and all other input nodes' default parent is this ancestor node.

2. Create a map to store name and TreeNode relationship for quick check.

class TreeNode {
    public String val;
    public TreeNode parent;
    public List<TreeNode> children;

    public TreeNode(String val) {
        this.val = val;
        children = new ArrayList<TreeNode>();
    }
}

import java.util.*;

public class Solution {
    public static void main(String[] args) {
        String s = "Frank->Mary,Mary->Sam,Mary->Bob,Sam->Katie,Sam->Pete";
        int target = 2;
        // Construct family tree

        TreeNode ancestor = new TreeNode("Ancestor");
        Map<String, TreeNode> map = new HashMap<String, TreeNode>();
        map.put("Ancestor", ancestor);

        String[] relations = s.split(",");

        for (String relation : relations) {
            String[] names = relation.split("->");
            String parent = names[0];
            String child = names[1];
            TreeNode parentNode, childNode;
            if (map.containsKey(parent)) {
                parentNode = map.get(parent);
            } else {
                parentNode = new TreeNode(parent);
                parentNode.parent = ancestor;
            }
            if (map.containsKey(child)) {
                childNode = map.get(child);
            } else {
                childNode = new TreeNode(child);
                childNode.parent = parentNode;
            }
            List<TreeNode> childrenList = parentNode.children;
            if (!childrenList.contains(childNode))
                childrenList.add(childNode);
            map.put(parent, parentNode);
            map.put(child, childNode);
            System.out.println(parent);
            System.out.println(child);

        }
        // Find children of ancestor
        List<TreeNode> firstChildren = ancestor.children;
        for (String tmp : map.keySet()) {
            TreeNode tmpNode = map.get(tmp);
            if (tmpNode.parent == ancestor) {
                firstChildren.add(tmpNode);
                System.out.println(tmpNode.val);
            }
        }

        System.out.println("start BFS");
        // BFS to get result
        int level = 0;
        List<TreeNode> currentList = new ArrayList<TreeNode>();
        List<String> result = new ArrayList<String>();
        List<TreeNode> nextList;
        currentList.add(ancestor);
        while (currentList.size() > 0) {
            nextList = new ArrayList<TreeNode>();
            for (TreeNode tmpNode : currentList) {
                List<TreeNode> childrenList = tmpNode.children;
                for (TreeNode oneChild : childrenList) {
                    if (!nextList.contains(oneChild))
                        nextList.add(oneChild);
                }
                currentList = nextList;
                level++;
                if (level == target) {
                    for (TreeNode tmpNode2 : currentList)
                        result.add(tmpNode2.val);
                    break;
                }
            }
        }
        for (String output : result) {
            System.out.println(output);
        }

    }
}