We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
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 ArrayList 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);
}
}
abstract class TreeVis {
public abstract int getResult();
public abstract void visitNode(TreeNode node);
public abstract void visitLeaf(TreeLeaf leaf);
}
class SumInLeavesVisitor extends TreeVis {
private int sumInLeaves = 0;
public int getResult() {
return sumInLeaves;
}
public void visitNode(TreeNode node) {
// Do nothing for nodes
}
public void visitLeaf(TreeLeaf leaf) {
sumInLeaves += leaf.getValue();
}
}
class ProductOfRedNodesVisitor extends TreeVis {
private long productOfRedNodes = 1L;
private static final int MOD = 1000000007;
public int getResult() {
return (int) productOfRedNodes;
}
private void multiply(Tree tree) {
if (tree.getColor() == Color.RED) {
productOfRedNodes = (productOfRedNodes * tree.getValue()) % MOD;
}
}
public void visitNode(TreeNode node) {
multiply(node);
}
public void visitLeaf(TreeLeaf leaf) {
multiply(leaf);
}
}
class FancyVisitor extends TreeVis {
private int sumOfValuesNonLeafEvenDepth = 0;
private int sumOfValuesGreenLeaf = 0;
public int getResult() {
return Math.abs(sumOfValuesNonLeafEvenDepth - sumOfValuesGreenLeaf);
}
public void visitNode(TreeNode node) {
if (node.getDepth() % 2 == 0) {
sumOfValuesNonLeafEvenDepth += node.getValue();
}
}
public void visitLeaf(TreeLeaf leaf) {
if (leaf.getColor() == Color.GREEN) {
sumOfValuesGreenLeaf += leaf.getValue();
}
}
}
public class Solution {
static Map tree = new HashMap<>();
public static Tree solve() {
Scanner sc = new Scanner(System.in);
int n = sc.nextInt();
Map<Integer, Object[]> nodeAtts = new HashMap<>();
// Read node values
for (int i = 0; i < n; i++) {
nodeAtts.put(i + 1, new Object[]{sc.nextInt(), null});
}
// Read node colors
for (int i = 0; i < n; i++) {
nodeAtts.get(i + 1)[1] = sc.nextInt() == 0 ? Color.RED : Color.GREEN;
}
// Read edges
Map<Integer, ArrayList<Integer>> edges = new HashMap<>();
for (int i = 1; i <= n; i++) {
edges.put(i, new ArrayList<>());
}
for (int i = 1; i < n; i++) {
int u = sc.nextInt();
int v = sc.nextInt();
edges.get(u).add(v);
edges.get(v).add(u);
}
// Create root node
Tree root = new TreeNode((Integer) nodeAtts.get(1)[0], (Color) nodeAtts.get(1)[1], 0);
tree.put(1, root);
// Perform DFS to build tree
DFS(n, edges, nodeAtts);
return tree.get(1);
}
private static void DFS(int n, Map<Integer, ArrayList<Integer>> edges, Map<Integer, Object[]> nodeAtts) {
boolean[] visited = new boolean[n + 1];
TreeNode parent = (TreeNode) tree.get(1);
DFSUtil(parent, 1, visited, edges, nodeAtts);
}
private static void DFSUtil(TreeNode parent, int v, boolean[] visited, Map<Integer, ArrayList<Integer>> edges, Map<Integer, Object[]> nodeAtts) {
visited[v] = true;
if (v != 1 && edges.get(v).size() == 1) {
TreeLeaf treeLeaf = new TreeLeaf((Integer) nodeAtts.get(v)[0], (Color) nodeAtts.get(v)[1], parent.getDepth() + 1);
parent.addChild(treeLeaf);
tree.put(v, treeLeaf);
return;
}
TreeNode treeNode = (v != 1) ? new TreeNode((Integer) nodeAtts.get(v)[0], (Color) nodeAtts.get(v)[1], parent.getDepth() + 1) : (TreeNode) tree.get(1);
if (v != 1) {
parent.addChild(treeNode);
tree.put(v, treeNode);
}
for (int n : edges.get(v)) {
if (!visited[n]) {
DFSUtil(treeNode, n, visited, edges, nodeAtts);
}
}
}
public static void main(String[] args) {
Tree root = solve();
SumInLeavesVisitor vis1 = new SumInLeavesVisitor();
ProductOfRedNodesVisitor vis2 = new ProductOfRedNodesVisitor();
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());
}
Java Visitor Pattern
You are viewing a single comment's thread. Return to all comments →
This code works fine for me
import java.util.ArrayList; import java.util.HashMap; import java.util.Map; import java.util.Scanner; import java.util.Iterator;
enum Color { RED, GREEN }
abstract class Tree { private int value; private Color color; private int depth;
}
class TreeNode extends Tree { private ArrayList children = new ArrayList<>();
}
class TreeLeaf extends Tree { public TreeLeaf(int value, Color color, int depth) { super(value, color, depth); }
}
abstract class TreeVis { public abstract int getResult(); public abstract void visitNode(TreeNode node); public abstract void visitLeaf(TreeLeaf leaf); }
class SumInLeavesVisitor extends TreeVis { private int sumInLeaves = 0;
}
class ProductOfRedNodesVisitor extends TreeVis { private long productOfRedNodes = 1L; private static final int MOD = 1000000007;
}
class FancyVisitor extends TreeVis { private int sumOfValuesNonLeafEvenDepth = 0; private int sumOfValuesGreenLeaf = 0;
}
public class Solution { static Map tree = new HashMap<>();
}