diff --git a/2020/day10.go b/2020/day10.go index 49ab39e..69934a1 100644 --- a/2020/day10.go +++ b/2020/day10.go @@ -19,8 +19,75 @@ func getMax(values map[int]bool) int { return max } +type Node struct { + Value int + Children []*Node + Paths int // the total number of paths from here to some target +} + +// Build a tree of possible paths, then count them. +func findArrangements(values map[int]bool, target int) int { + root := buildTree(values, target) + return countPaths(root, target) +} + +func buildTree(values map[int]bool, target int) *Node { + values[0] = true + nodeMap := make(map[int]*Node) + + for i := 0; i < target; i++ { + if values[i] { + node := createOrFetchNode(nodeMap, i) + + for j := 1; j < 4; j++ { + if i+j > target { + break + } + + if values[i+j] { + child := createOrFetchNode(nodeMap, i+j) + node.Children = append(node.Children, child) + } + } + } + } + + return nodeMap[0] +} + +func createOrFetchNode(nodeMap map[int]*Node, value int) *Node { + if node, ok := nodeMap[value]; ok { + return node + } + + node := &Node{ + Value: value, + Children: make([]*Node, 0), + } + nodeMap[value] = node + + return node +} + +func countPaths(node *Node, target int) int { + if node.Value == target { + return 1 + } + + paths := 0 + for _, child := range node.Children { + if child.Paths != 0 { + paths += child.Paths + } else { + paths += countPaths(child, target) + } + } + node.Paths = paths + return paths +} + func main() { - // step := os.Args[1] + step := os.Args[1] values := fileutils.InputParserIntMap(os.Args[2]) diffMap := make(map[int]int) @@ -28,21 +95,27 @@ func main() { device := getMax(values) + 3 values[device] = true - diff := 0 - for i := 0; i < device; i++ { - // increment diff, make sure we haven't broken the chain - // Note that we're actually logically checking the joltage at i+1 - // but that serves us well here - diff++ - if diff > 3 { - log.Panicf("Diff too big, bailing.") - } + switch step { + case "1": + diff := 0 + for i := 0; i < device; i++ { + // increment diff, make sure we haven't broken the chain + // Note that we're actually logically checking the joltage at i+1 + // but that serves us well here + diff++ + if diff > 3 { + log.Panicf("Diff too big, bailing.") + } - // if we have a device at this joltage, register and reset count - if values[i+1] { - diffMap[diff]++ - diff = 0 + // if we have a device at this joltage, register and reset count + if values[i+1] { + diffMap[diff]++ + diff = 0 + } } + fmt.Println(diffMap[1] * diffMap[3]) + + case "2": + fmt.Println(findArrangements(values, device)) } - fmt.Println(diffMap[1] * diffMap[3]) }