aoc2021/twelve/navigation.go

package twelve

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

type caveSize int

const (
	startLabel          = "start"
	endLabel            = "end"
	small      caveSize = 0
	large      caveSize = 1
)

type cave struct {
	label       string
	connections []*cave
	size        caveSize
}

type navigation struct {
	input []string
	start *cave
	end   *cave
	caves map[string]*cave
}

func (n *navigation) load(filename string) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	n.caves = map[string]*cave{}
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		n.input = append(n.input, scanner.Text())
		caves := strings.Split(scanner.Text(), "-")

		if len(caves) != 2 {
			return fmt.Errorf("expected 2 caves connected, found %s", scanner.Text())
		}

		if n.caves[caves[0]] == nil {
			n.caves[caves[0]] = &cave{
				label: caves[0],
				size:  size(caves[0]),
			}
		}

		if n.caves[caves[1]] == nil {
			n.caves[caves[1]] = &cave{
				label: caves[1],
				size:  size(caves[1]),
			}
		}

		// Add connection to 1 on 0
		found := false
		for _, c := range n.caves[caves[0]].connections {
			if c.label == caves[1] {
				found = true
			}
		}
		if !found {
			n.caves[caves[0]].connections = append(n.caves[caves[0]].connections, n.caves[caves[1]])
		}

		// Add connection to 0 on 1
		found = false
		for _, c := range n.caves[caves[1]].connections {
			if c.label == caves[0] {
				found = true
			}
		}
		if !found && caves[0] != startLabel {
			n.caves[caves[1]].connections = append(n.caves[caves[1]].connections, n.caves[caves[0]])
		}
	}
	// Initialize start/end
	n.findStart()
	n.findEnd()

	return nil
}

func size(label string) caveSize {
	if strings.ToUpper(label) == label {
		return large
	}

	return small
}

func (n *navigation) findStart() *cave {
	if n.start != nil {
		return n.start
	}

	for k := range n.caves {
		if k == startLabel {
			n.start = n.caves[k]
		}
	}

	return n.start
}

func (n *navigation) findEnd() *cave {
	if n.end != nil {
		return n.end
	}

	for k := range n.caves {
		if k == endLabel {
			n.end = n.caves[k]
		}
	}

	return n.end
}

func (n *navigation) findPaths() [][]string {
	s := n.findStart()
	foundPaths := [][]string{}

	for _, c := range s.connections {
		paths := n.findPathsRecursive([]string{s.label}, c, 1)
		for _, p := range paths {
			labels := []string{"start"}
			for _, c := range p {
				labels = append(labels, c.label)
			}
			foundPaths = append(foundPaths, labels)
		}
	}

	return foundPaths
}

func (n *navigation) findPaths2() [][]string {
	s := n.findStart()
	foundPaths := [][]string{}

	for _, c := range s.connections {
		paths := n.findPathsRecursive([]string{s.label}, c, 2)
		for _, p := range paths {
			labels := []string{"start"}
			for _, c := range p {
				labels = append(labels, c.label)
			}
			foundPaths = append(foundPaths, labels)
		}
	}

	// for k := range n.caves {
	// 	fmt.Println(n.caves[k])
	// }

	// for _, p := range foundPaths {
	// 	fmt.Println(p)
	// }

	return foundPaths
}

func (n *navigation) findPathsRecursive(from []string, in *cave, smallCaveLimit int) [][]*cave {
	paths := [][]*cave{}

	if in == n.end {
		paths = append(paths, []*cave{
			n.end,
		})
		return paths
	}

	for _, c := range in.connections {
		if c.label == n.start.label {
			continue
		}

		// No back and forth shenanigans with small caves!
		// Check if we have already been to the small cave
		if c.size == small {
			connectionOption := true
			smallVisits := map[string]int{}
			for _, l := range from {
				if size(l) == small {
					smallVisits[l]++
				}
			}

			if in.size == small {
				smallVisits[in.label]++
			}

			if smallCaveLimit > 1 {
				for k := range smallVisits {
					if smallVisits[k] == smallCaveLimit {
						connectionOption = false
					}
				}
			}

			if smallVisits[c.label] == smallCaveLimit {
				connectionOption = false
			}

			if smallVisits[c.label] == 0 {
				connectionOption = true
			}

			if !connectionOption {
				continue
			}
		}

		pathsFromHere := n.findPathsRecursive(append(from, in.label), c, smallCaveLimit)
		if len(pathsFromHere) == 0 {
			continue
		}

		for _, path := range pathsFromHere {
			// Add the current to the front of each new path
			newOption := []*cave{
				in,
			}
			newOption = append(newOption, path...)

			paths = append(paths, newOption)
		}
	}

	return paths
}