114 lines
1.9 KiB
Go
114 lines
1.9 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"regexp"
|
|
"strconv"
|
|
|
|
"git.annabunch.es/annabunches/adventofcode/2020/lib/util"
|
|
)
|
|
|
|
type Node struct {
|
|
Color string
|
|
Children map[*Node]int
|
|
}
|
|
|
|
var nodeMap map[string]*Node
|
|
|
|
var colorRe = regexp.MustCompile("^(.*?) bags contain")
|
|
var childRe = regexp.MustCompile("(?:contain|,) (\\d+) (.*?) bag")
|
|
|
|
func findOrCreateNode(color string) *Node {
|
|
if node, ok := nodeMap[color]; ok {
|
|
return node
|
|
}
|
|
|
|
node := &Node{
|
|
Color: color,
|
|
}
|
|
nodeMap[color] = node
|
|
return node
|
|
}
|
|
|
|
func parseRule(line string) {
|
|
if line == "" {
|
|
return
|
|
}
|
|
|
|
color := colorRe.FindStringSubmatch(line)[1]
|
|
|
|
node := findOrCreateNode(color)
|
|
|
|
children := make(map[*Node]int)
|
|
childrenData := childRe.FindAllStringSubmatch(line, 128)
|
|
|
|
for _, childData := range childrenData {
|
|
child := findOrCreateNode(childData[2])
|
|
count, err := strconv.Atoi(childData[1])
|
|
if err != nil {
|
|
log.Panicf(err.Error())
|
|
}
|
|
children[child] = count
|
|
}
|
|
|
|
node.Children = children
|
|
}
|
|
|
|
func findInDescendants(node *Node, color string) bool {
|
|
for child, _ := range node.Children {
|
|
if rFindInDescendants(child, color) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func rFindInDescendants(node *Node, color string) bool {
|
|
// found it
|
|
if node.Color == color {
|
|
return true
|
|
}
|
|
|
|
for child, _ := range node.Children {
|
|
if rFindInDescendants(child, color) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func countChildren(node *Node) int {
|
|
total := 1 // count ourself
|
|
for child, count := range node.Children {
|
|
total += count * countChildren(child)
|
|
}
|
|
return total
|
|
}
|
|
|
|
func main() {
|
|
step := os.Args[1]
|
|
values := util.InputParserStrings(os.Args[2])
|
|
|
|
nodeMap = make(map[string]*Node)
|
|
|
|
for _, line := range values {
|
|
parseRule(line)
|
|
}
|
|
|
|
switch step {
|
|
case "1":
|
|
total := 0
|
|
for _, node := range nodeMap {
|
|
if findInDescendants(node, "shiny gold") {
|
|
total++
|
|
}
|
|
}
|
|
fmt.Println("Total:", total)
|
|
|
|
case "2":
|
|
fmt.Println("Total:", countChildren(nodeMap["shiny gold"])-1)
|
|
}
|
|
}
|