_3004.java

packagecom.fishercoder.solutions.fourththousand;

importjava.util.ArrayList;
importjava.util.HashMap;
importjava.util.List;
importjava.util.Map;

publicclass_3004 {
publicstaticclassSolution1 {
/*
 * My completely original solution.
 * Practice makes perfect!
 * Post-order traversal is the way to go since we need to process all children first before processing any particular node.
 */
classColoredTreeNode {
intval;
intcolor;
List<ColoredTreeNode> children;
booleanallSubtreeSameColor;
inttotalChildrenCount;

publicColoredTreeNode(intval, intcolor) {
this.val = val;
this.color = color;
this.children = newArrayList<>();
this.allSubtreeSameColor =
true; // initialize to be true until it's built/proven to be false
this.totalChildrenCount = 1; // count itself as its own child
 }
 }

intmaxSize = 0;

publicintmaximumSubtreeSize(int[][] edges, int[] colors) {
if (edges == null || edges.length == 0 || edges[0].length == 0) {
returncolors.length > 0 ? 1 : 0;
 }
ColoredTreeNoderoot = buildTree(edges, colors);
inttotalNodeCount = postOrder(root);
if (root.allSubtreeSameColor) {
returntotalNodeCount;
 }
returnmaxSize;
 }

privateintpostOrder(ColoredTreeNoderoot) {
if (root == null) {
return0;
 }
inttotalChildrenCount = 1; // count itself as a child
for (ColoredTreeNodechild : root.children) {
intcount = postOrder(child);
totalChildrenCount += count;
if (root.color != child.color || !child.allSubtreeSameColor) {
root.allSubtreeSameColor = false;
 }
 }
root.totalChildrenCount = totalChildrenCount;
if (root.allSubtreeSameColor) {
maxSize = Math.max(maxSize, root.totalChildrenCount);
 }
returntotalChildrenCount;
 }

privateColoredTreeNodebuildTree(int[][] edges, int[] colors) {
Map<Integer, ColoredTreeNode> map = newHashMap<>();
for (inti = 0; i < edges.length; i++) {
ColoredTreeNodeparent =
map.getOrDefault(
edges[i][0], newColoredTreeNode(edges[i][0], colors[edges[i][0]]));
ColoredTreeNodechild =
map.getOrDefault(
edges[i][1], newColoredTreeNode(edges[i][1], colors[edges[i][1]]));
parent.children.add(child);
map.put(edges[i][0], parent);
map.put(edges[i][1], child);
 }
returnmap.get(0);
 }
 }
}