Java Visitor Pattern

  • + 0 comments

    java solution with the supposed pre-existing classes

    public class Solution {
        static int n;
        static int[] values;
        static Color[] colors;
        static Map<Integer, Set<Integer>> edges;
    
        public static void main(String[] args) {
            Tree root = solve();
    
            SumInLeavesVisitor vis1 = new SumInLeavesVisitor();
            ProductRedNodesVisitor vis2 = new ProductRedNodesVisitor();
            FancyVisitor vis3 = new FancyVisitor();
    
            root.accept(vis1);
            root.accept(vis2);
            root.accept(vis3);
    
            System.out.println(vis1.getResult());
            System.out.println(vis2.getResult());
            System.out.println(vis3.getResult());
        }
        
         public static Tree solve() {
            Scanner sc = new Scanner(System.in);
            n = sc.nextInt();
    
            values = new int[n+1];
            for (int i = 1; i <= n; i++) {
                values[i] = sc.nextInt();
            }
    
            colors = new Color[n+1];
            for (int i = 1; i <= n; i++) {
                int c = sc.nextInt();
                colors[i] = (c == 0 ? Color.RED : Color.GREEN);
            }
    
            edges = new HashMap<>();
            for (int i = 1; i <= n; i++) edges.put(i, new HashSet<>());
            for (int i = 0; i < n-1; i++) {
                int u = sc.nextInt();
                int v = sc.nextInt();
                edges.get(u).add(v);
                edges.get(v).add(u);
            }
    
            sc.close();
            return buildTree(1, 0, -1); // root=1, depth=0, parent=-1
        }    
        
        private static Tree buildTree(int node, int depth, int parent) {
            Set<Integer> childs = edges.get(node);
            if(childs.size() == 1 && childs.contains(parent)){
                //es hoja
                return new TreeLeaf(values[node], colors[node], depth);
            } else {
                TreeNode tn = new TreeNode(values[node], colors[node], depth);
                for(int child : childs){
                    if(child != parent) {
                        tn.addChild(buildTree(child, depth+1, node));
                    }
                }
                return tn;
            }
        }
    }
    
    enum Color {RED, GREEN}
    
    abstract class Tree {
        private int value;
        private Color color;
        private int depth;
        
        public Tree(int value, Color color, int depth){
            this.value = value;
            this.color = color;
            this.depth = depth;
            
        }
        
        public int getValue(){ return value;}
        public Color getColor(){ return color;}
        public int getDepth(){ return depth;}
        
        public abstract void accept(TreeVis visitor);
    }
    
    class TreeNode extends Tree{
        private List<Tree> children = new ArrayList<>();
        
        public TreeNode(int value, Color color, int depth){
            super(value, color, depth);
        }
        
        public void accept(TreeVis visitor){
            visitor.visitNode(this);
            for(Tree child : children){
                child.accept(visitor);
            }
        }
        
        public void addChild(Tree child){ children.add(child); }
    }
    
    class TreeLeaf extends Tree{
        public TreeLeaf(int value, Color color, int depth){
            super(value, color, depth);
        }
        
        public void accept(TreeVis visitor){
            visitor.visitLeaf(this);
        }
    }
    
    interface TreeVis{
        int getResult();
        void visitNode(TreeNode tree);
        void visitLeaf(TreeLeaf tree);
    }
    
    class SumInLeavesVisitor implements TreeVis {
        private int sum = 0;
        
        public int getResult(){ return sum; }
        
        public void visitNode(TreeNode node){ /*no hace nada*/ }
        
        public void visitLeaf(TreeLeaf leaf){
            sum += leaf.getValue();
        }
    }
    
    class ProductRedNodesVisitor implements TreeVis {
        private long product = 1;
        private static final int MOD = 1000000007;
         
        public int getResult(){ return (int)(product % MOD); }
        
        public void visitNode(TreeNode node){ 
            if(node.getColor() == Color.RED)
                product = (product * node.getValue()) % MOD;
         }
        
        public void visitLeaf(TreeLeaf leaf) {
            if (leaf.getColor() == Color.RED) {
                product = (product * leaf.getValue()) % MOD;
            }
        }
    }
    
    class FancyVisitor  implements TreeVis {
        private int evenDepthNonLeafSum  = 0;
        private int greenLeafSum = 0;
         
        public int getResult(){ return (int)(Math.abs( evenDepthNonLeafSum - greenLeafSum)); }
        
        public void visitNode(TreeNode node){ 
            if(node.getDepth()  % 2 == 0)
                evenDepthNonLeafSum += node.getValue();
         }
        
        public void visitLeaf(TreeLeaf leaf) {
            if (leaf.getColor() == Color.GREEN) {
                greenLeafSum += leaf.getValue();
            }
        }
    }