adventofcode/2017/day07/day07.go

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
}