@@ -56,38 +56,38 @@ private Map<String, Program> buildTree(String[] lines) {
5656 }
5757
5858 private int findCorrectWeight (Program root ) {
59- // For each node, check if its children are balanced
60- for (Program program : root .children ) {
61- int result = findCorrectWeight (program );
62- if (result != -1 ) return result ;
59+ // Check if any children are unbalanced
60+ for (Program child : root .children ) {
61+ int result = findCorrectWeight (child );
62+ if (result != -1 ) {
63+ return result ;
64+ }
6365 }
6466
65- // Get weights of all siblings
66- if (root .parent != null ) {
67+ // Check if this node's children are balanced
68+ if (! root .children . isEmpty () ) {
6769 Map <Integer , List <Program >> weightGroups = new HashMap <>();
68- for (Program sibling : root .parent .children ) {
69- int weight = sibling .getTotalWeight ();
70- weightGroups .computeIfAbsent (weight , k -> new ArrayList <>()).add (sibling );
70+ for (Program child : root .children ) {
71+ weightGroups .computeIfAbsent (child .getTotalWeight (), k -> new ArrayList <>()).add (child );
7172 }
7273
73- // If we have exactly two different weights and one is unique
7474 if (weightGroups .size () == 2 ) {
75- int correctWeight = - 1 ;
76- int wrongWeight = - 1 ;
77- Program wrongProgram = null ;
75+ // Find the odd one out
76+ Map . Entry < Integer , List < Program >> wrongEntry = null ;
77+ Map . Entry < Integer , List < Program >> correctEntry = null ;
7878
7979 for (Map .Entry <Integer , List <Program >> entry : weightGroups .entrySet ()) {
8080 if (entry .getValue ().size () == 1 ) {
81- wrongWeight = entry .getKey ();
82- wrongProgram = entry .getValue ().get (0 );
81+ wrongEntry = entry ;
8382 } else {
84- correctWeight = entry . getKey () ;
83+ correctEntry = entry ;
8584 }
8685 }
8786
88- if (wrongProgram != null ) {
89- int weightDiff = wrongWeight - correctWeight ;
90- return wrongProgram .weight - weightDiff ;
87+ if (wrongEntry != null && correctEntry != null ) {
88+ Program wrongProgram = wrongEntry .getValue ().get (0 );
89+ int difference = wrongEntry .getKey () - correctEntry .getKey ();
90+ return wrongProgram .weight - difference ;
9191 }
9292 }
9393 }
0 commit comments