Solve day 11

day11/main.go

@@ -1,7 +1,7 @@
 package main
 
 import (
-	"errors"
+	"cmp"
 	"fmt"
 	"io"
 	"os"
@@ -35,36 +35,119 @@ func main() {
 
 	input := strings.TrimSpace(string(inputBytes))
 	inputLines := strings.Split(input, "\n")
-	expandedInputLines, err := expandInput(inputLines)
+
+	nodes, err := parseInputMatrix(inputLines)
 	if err != nil {
-		panic(fmt.Sprintf("failed to expand input: %s", err))
+		panic(fmt.Sprintf("failed to parse input: %s", err))
 	}
 
-	nodes := parseInputMatrix(expandedInputLines)
 	fmt.Printf("Part 1: %d\n", part1(nodes))
+	fmt.Printf("Part 2: %d\n", part2(nodes))
 }
 
 func part1(nodes []Coordinate) int {
+	expanded := expandUniverse(nodes, 2)
+	return computePairwiseDistanceTotal(expanded)
+}
+
+func part2(nodes []Coordinate) int {
+	expanded := expandUniverse(nodes, 1_000_000)
+	return computePairwiseDistanceTotal(expanded)
+}
+
+func computePairwiseDistanceTotal(nodes []Coordinate) int {
 	type Pair struct {
 		node1 Coordinate
 		node2 Coordinate
 	}
 
-	pairs := []Pair{}
+	total := 0
 	for i := 0; i < len(nodes)-1; i++ {
 		for j := i + 1; j < len(nodes); j++ {
 			pair := Pair{node1: nodes[i], node2: nodes[j]}
-			pairs = append(pairs, pair)
+			distance := abs(pair.node2.col-pair.node1.col) + abs(pair.node2.row-pair.node1.row)
+
+			total += distance
 		}
 	}
 
-	total := 0
-	for _, pair := range pairs {
-		distance := abs(pair.node2.col-pair.node1.col) + abs(pair.node2.row-pair.node1.row)
-		total += abs(distance)
+	return total
+}
+
+func expandUniverse(nodes []Coordinate, expansion int) []Coordinate {
+	rowExpanded := expandAlongAxis(
+		nodes,
+		expansion,
+		func(c Coordinate) int { return c.row },
+		func(c Coordinate, n int) Coordinate { return Coordinate{row: n, col: c.col} },
+	)
+
+	expanded := expandAlongAxis(
+		rowExpanded,
+		expansion,
+		func(c Coordinate) int { return c.col },
+		func(c Coordinate, n int) Coordinate { return Coordinate{row: c.row, col: n} },
+	)
+
+	return expanded
+}
+
+// expandAlongAxis will expand by the given amount the universe only along a single axis.
+// getAxis will allow the function to get the value of a single axis, given a coordinate,
+// and setAxis must return a new coordinate with the same axis set to the given value.
+func expandAlongAxis(
+	nodes []Coordinate,
+	expansion int,
+	getAxis func(Coordinate) int,
+	setAxisValue func(Coordinate, int) Coordinate,
+) []Coordinate {
+	expanded := slices.Clone(nodes)
+	slices.SortFunc(expanded, func(a, b Coordinate) int { return cmp.Compare(getAxis(a), getAxis(b)) })
+
+	blank := findBlankAxisValues(nodes, getAxis)
+	totalExpansion := 0
+	for _, blankIdx := range blank {
+		for i := range expanded {
+			axisValue := getAxis(expanded[i])
+			if axisValue <= blankIdx+totalExpansion {
+				continue
+			}
+
+			expanded[i] = setAxisValue(expanded[i], axisValue+(expansion-1))
+		}
+		totalExpansion += (expansion - 1)
 	}
 
-	return total
+	return expanded
+}
+
+// findBlankAxisValues finds all the positions where all values along the given axis are blank
+func findBlankAxisValues(nodes []Coordinate, getAxis func(Coordinate) int) []int {
+	maxCoord := slices.MaxFunc(nodes, func(node1, node2 Coordinate) int {
+		return cmp.Compare(getAxis(node1), getAxis(node2))
+	})
+
+	axisMax := getAxis(maxCoord)
+	knownAlongAxis := mapToSet(nodes, getAxis)
+	blank := []int{}
+	for i := 0; i <= axisMax; i++ {
+		if _, ok := knownAlongAxis[i]; !ok {
+			blank = append(blank, i)
+		}
+	}
+
+	return blank
+}
+
+// mapToSet maps all values of the given input and turns them into a set
+func mapToSet[T any, U comparable, S ~[]T](items S, mapper func(T) U) map[U]struct{} {
+	res := map[U]struct{}{}
+	for _, item := range items {
+		key := mapper(item)
+		res[key] = struct{}{}
+	}
+
+	return res
 }
 
 func abs(x int) int {
@@ -75,93 +158,18 @@ func abs(x int) int {
 	return x
 }
 
-func expandInput(inputLines []string) ([][]rune, error) {
-	if len(inputLines) == 0 {
-		return nil, nil
-	}
-
-	runeMatrix := makeRuneMatrix(inputLines)
-	expanded1 := expandInputVertically(runeMatrix)
-	transposed1, err := transposeMatrix(expanded1)
-	if err != nil {
-		return nil, err
-	}
-
-	expanded2 := expandInputVertically(transposed1)
-	transposed2, err := transposeMatrix(expanded2)
-	if err != nil {
-		return nil, err
-	}
-
-	return transposed2, nil
-}
-
-func makeRuneMatrix(strings []string) [][]rune {
-	matrix := make([][]rune, len(strings))
-	for i, line := range strings {
-		matrix[i] = make([]rune, len(line))
-		for j, char := range line {
-			matrix[i][j] = char
-		}
-	}
-
-	return matrix
-}
-
-func expandInputVertically(inputMatrix [][]rune) [][]rune {
-	expanded := [][]rune{}
-	for _, line := range inputMatrix {
-		expanded = append(expanded, line)
-		if count(line, '.') == len(line) {
-			expanded = append(expanded, slices.Clone(line))
-		}
-	}
-
-	return expanded
-}
-
-func transposeMatrix[T any, S ~[]T, M []S](matrix M) (M, error) {
-	for _, row := range matrix {
-		if len(row) != len(matrix[0]) {
-			return nil, errors.New("not all rows are the same length")
-		}
-	}
-
-	transposed := make(M, 0)
-	for range matrix {
-		transposed = append(transposed, make(S, len(matrix)))
-	}
-
-	for i, row := range matrix {
-		for j, item := range row {
-			transposed[j][i] = item
-		}
-	}
-
-	return transposed, nil
-}
-
-func count[T comparable, S ~[]T](s S, target T) int {
-	count := 0
-	for _, item := range s {
-		if item == target {
-			count++
-		}
-	}
-
-	return count
-}
-
-func parseInputMatrix(input [][]rune) []Coordinate {
+func parseInputMatrix(input []string) ([]Coordinate, error) {
 	nodes := []Coordinate{}
 
 	for row, line := range input {
 		for col, char := range line {
 			if char == '#' {
 				nodes = append(nodes, Coordinate{row: row, col: col})
+			} else if char != '.' {
+				return nil, fmt.Errorf("invalid input char %c", char)
 			}
 		}
 	}
 
-	return nodes
+	return nodes, nil
 }