Java Visitor Pattern

  • + 0 comments
    class SumInLeavesVisitor extends TreeVis {
        private int sum = 0;
    
        public int getResult() {
            return sum;
        }
    
        public void visitNode(TreeNode node) {
            // No action needed for internal nodes in this visitor
        }
    
        public void visitLeaf(TreeLeaf leaf) {
            sum += leaf.getValue();
        }
    }
    
    class ProductOfRedNodesVisitor extends TreeVis {
        private double product = 1;
    
        public int getResult() {
            return (int) product;
        }
    
        public void visitNode(TreeNode node) {
            increaseProduct(node);
        }
    
        public void visitLeaf(TreeLeaf leaf) {
            increaseProduct(leaf);
        }
    
        private void increaseProduct(Tree tree) {
            if (tree.getColor() == Color.RED) {
                product = (product * tree.getValue()) % (Math.pow(10, 9) + 7);
            }
        }
    }
    
    class FancyVisitor extends TreeVis {
        private int nonLeafNodeSum = 0;
        private int greenLeafSum = 0;
    
        public int getResult() {
            return Math.abs(nonLeafNodeSum - greenLeafSum);
        }
    
        public void visitNode(TreeNode node) {
            if (node.getDepth() % 2 == 0) {
                nonLeafNodeSum += node.getValue();
            }
        }
    
        public void visitLeaf(TreeLeaf leaf) {
            if (leaf.getColor() == Color.GREEN) {
                greenLeafSum += leaf.getValue();
            }
        }
    }
    
    public class Solution {
      
        public static Tree solve() {
            Scanner scan = new Scanner(System.in);
            int n = Integer.parseInt(scan.nextLine());
            String valuesLine = scan.nextLine();
            String colorsLine = scan.nextLine();
    
            String[] valuesStr = valuesLine.split(" ");
            int[] values = new int[n];
    
            for (int i = 0; i < n; i++) {
                values[i] = Integer.parseInt(valuesStr[i]);
            }
    
            String[] colorsStr = colorsLine.split(" ");
            byte[] colors = new byte[n];
            
            for (int i = 0; i < n; i++) {
                colors[i] = (byte) Integer.parseInt(colorsStr[i]);
            }
    
            Map<Integer, List<Integer>> adj = new HashMap<>();
            boolean[] visited = new boolean[n];
    
            for (int i = 0; i < n; i++) {
                adj.put(i, new ArrayList<Integer>());
            }
    
            while (scan.hasNextInt()) {
                int u = scan.nextInt() - 1;
                int v = scan.nextInt() - 1;
                adj.get(u).add(v);
                adj.get(v).add(u);
            }
    
            scan.close();
            
            return buildTree(0, 0, values, colors, adj, visited);
        }
    
        public static Tree buildTree(
            int index, int depth, int[] values, byte[] colors, Map<Integer, List<Integer>> adj, boolean[] visited
        ) {
            boolean isLeaf = true;
            visited[index] = true;
    
            Color color = colors[index] == 0 ? Color.RED : Color.GREEN;
            List<Integer> children = adj.get(index);
            TreeNode treeNode = new TreeNode(values[index], color, depth);
    
            for (int child : children) {
                if (!visited[child]) {
                    isLeaf = false;
                    Tree childNode = buildTree(child, depth + 1, values, colors, adj, visited);
                    treeNode.addChild(childNode);
                }
            }
    
            if (isLeaf) {
                return new TreeLeaf(values[index], color, depth);
            } else {
                return treeNode;
            }
        }