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;
        }
    }