Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- public static Tree solve() {
- //read the tree from STDIN and return its root as a return value of this function
- //List to contain all nodes/leaves
- List<Tree> trees = new ArrayList<>();
- try (BufferedReader reader = new BufferedReader(new InputStreamReader(System.in))) {
- //read in the amount of nodes/leaves
- int t = Integer.parseInt(reader.readLine());
- //debug
- System.out.println(t);
- //arrays for containing all the values/colours of each node/leaf; there will be 't' amount of values
- int[] values = new int[t];
- int[] colors = new int[t];
- //read the values/colours and then split them
- String[] valueInput = reader.readLine().split("");
- String[] colorInout = reader.readLine().split("");
- //convert the arrays of Strings of colours/values to int values and store them in the int[] arrays
- for (int i = 0; i < t; i++) {
- values[i] = Integer.parseInt(valueInput[i]);
- colors[i] = Integer.parseInt(colorInout[i]);
- }
- //debug
- System.out.println(values);
- System.out.println(colors);
- //read in the edges and store them in a List
- List<String> edgeInputs = new ArrayList<>();
- String edge;
- while(!(edge = reader.readLine()).equals("")) {
- edgeInputs.add(edge);
- }
- //split each edge into its component parts, store the parent in one List, and store the child in another List
- //both should be at the same index in their relevant List
- //store both Lists in a List<List<String>>
- List<List<Integer>> edges = new ArrayList<>();
- edges.add(new ArrayList<>());
- edges.add(new ArrayList<>());
- edgeInputs.forEach(s -> {
- String[] nodes = s.split(" ");
- edges.get(0).add(Integer.parseInt(nodes[0]));
- edges.get(1).add(Integer.parseInt(nodes[1]));
- });
- //checks if the root is a node or leaf, creates the correct tree object, and adds it to the List
- if (isLeaf(1, edges)) {
- trees.add(new TreeLeaf(values[0], color(colors[0]), 0));
- } else {
- trees.add(new TreeNode(values[0], color(colors[0]), 0));
- }
- //iterates over the remaining values
- //determines if they're nodes or leaves, and creates a corresponding Tree object
- //finds the parent from the List of created Trees, adds the new Tree to the parent's List of children
- //adds the new Tree to the List
- for (int i = 1; i < values.length; i++) {
- Tree tree;
- int index = 1 + 1;
- TreeNode parent = (TreeNode) trees.get(getParentIndex(index, edges));
- int depth = parent.getDepth() + 1;
- if (isLeaf(index, edges)) {
- tree = new TreeLeaf(values[i], color(colors[i]), depth);
- } else {
- tree = new TreeNode(values[i], color(colors[i]), depth);
- }
- parent.addChild(tree);
- trees.add(tree);
- }
- } catch (Exception e) {}
- return trees.get(0);
- }
- //checks if a Tree is a leaf
- //searches for its index in the List containing parents
- //if it doesn't appear, then it must be a leaf, not a node, return true
- public static boolean isLeaf(int index, List<List<Integer>> edges) {
- return !edges.get(0).contains(index);
- }
- //generates the Color based on the value
- //0 == Red / 1 == Green
- public static Color color(int color) {
- return color == 0 ? Color.RED : Color.GREEN;
- }
- //finds where a node/leaf's parent is in the List containing parents
- //by searching for
- public static int getParentIndex(int index, List<List<Integer>> edges) {
- int i = edges.get(1).indexOf(index);
- return edges.get(0).get(i) - 1;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement