package main import ( "bufio" "fmt" "log" "os" "strconv" "strings" ) type Node struct { Name string Parent string Children []string Num int } // A list of all nodes we've added var track map[string]*Node var root *Node func main() { track = make(map[string]*Node) inp := StdinToStrings() totalWeight := 0 for i := range inp { n := parseNode(inp[i]) totalWeight += n.Num updateNodeInTrack(n) } fmt.Println("=== Results ===") n := findRoot() fmt.Println("Root Node: " + n.Name) fmt.Println("Total Weight:", totalWeight, "\n") imb, _ := findImbalance(n.Name) imbChain := []string{n.Name, imb} fmt.Println("") for imb != "" { fmt.Println("Imbalance on", imb) imb, _ = findImbalance(imb) fmt.Println("") if imb != "" { imbChain = append(imbChain, imb) } } if len(imbChain) > 2 { // imbChain[len(imbChain)-3] is the parent parent := imbChain[len(imbChain)-2] // imbChain[len(imbChain)-2] should be the one that needs balancing target := imbChain[len(imbChain)-1] fmt.Println("Balance At:", target) _, propWeight := findImbalance(parent) currWeight := getWeight(target) fmt.Println("\nAdjust weight of", target, "to:", (track[target].Num + (propWeight - currWeight))) } else { fmt.Println("Couldn't find an imbalance") } } // findImbalance returns the wrongly balanced node name // along with the weight that it _should_ be func findImbalance(nm string) (string, int) { fmt.Println(nm, ">>") weightTracker := make(map[int]int) chldWeights := getChildWeights(nm) for k, v := range chldWeights { fmt.Println(" ", k, " (", track[k].Num, ") :", v) weightTracker[v]++ } // Now find the different one var properWeight int var wrongNode string for k, v := range weightTracker { if v == 1 { for chK, chV := range chldWeights { if chV == k { wrongNode = chK } } } else { properWeight = k } } return wrongNode, properWeight } func getChildWeights(nm string) map[string]int { ret := make(map[string]int) for k := range track { if strings.HasSuffix(getLineage(k), nm+" < "+k) { ret[k] = getWeight(k) } } return ret } func getWeight(nm string) int { var weight int for i := range track { if strings.Contains(getLineage(i), nm+" <") { weight += track[i].Num } } return weight + track[nm].Num } func getLineage(nm string) string { ret := "" if i, ok := track[nm]; ok { ret = getLineage(i.Parent) + " < " + nm } return ret } func findRoot() *Node { for i := range track { if track[i].Parent == "" { return track[i] } } return nil } func updateNodeInTrack(n *Node) *Node { if _, ok := track[n.Name]; !ok { // We need to add it track[n.Name] = n } // Now update all child nodes already on the track with this as their parent for j := range n.Children { updateNodeInTrack(&Node{Name: n.Children[j], Parent: n.Name}) } // Now update the other values, if we've got 'em if n.Num > 0 { track[n.Name].Num = n.Num } if n.Parent != "" { track[n.Name].Parent = n.Parent } return track[n.Name] } func parseNode(inp string) *Node { n := new(Node) pts := strings.Split(inp, " ") for i := range pts { switch i { case 0: n.Name = pts[i] case 1: n.Num = Atoi(pts[i][1 : len(pts[i])-1]) case 2: // The arrow continue default: // All children childName := pts[i] if strings.HasSuffix(childName, ",") { childName = childName[:len(childName)-1] } n.Children = append(n.Children, childName) } } return n } func StdinToStrings() []string { var input []string scanner := bufio.NewScanner(os.Stdin) for scanner.Scan() { input = append(input, scanner.Text()) } return input } func Atoi(i string) int { var ret int var err error if ret, err = strconv.Atoi(i); err != nil { log.Fatal("Invalid Atoi") } return ret }