shape_solver.go

package main

import (
	"fmt"
	"sort"
	"strings"
)

type Shape struct {
	grid   [][]bool
	points []Point
}

func (s *Shape) Width() int {
	return len(s.grid[0])
}

func (s *Shape) Height() int {
	return len(s.grid)
}

func (s *Shape) ComputePoints() {
	for y := 0; y < s.Height(); y++ {
		for x := 0; x < s.Width(); x++ {
			if s.grid[y][x] {
				s.points = append(s.points, Point{x, y})
			}
		}
	}
}

func ShapeFromString(input string) *Shape {
	lines := strings.Split(input, "\n")
	height := len(lines)
	width := len(lines[0])
	shape := &Shape{}
	for y := 0; y < height; y++ {
		line := lines[y]
		var row []bool
		for x := 0; x < width; x++ {
			row = append(row, line[x] != ' ')
		}
		shape.grid = append(shape.grid, row)
	}

	return shape
}

func (s *Shape) String() string {
	var rows []string
	for y := 0; y < s.Height(); y++ {
		var row strings.Builder
		for x := 0; x < s.Width(); x++ {
			if s.grid[y][x] {
				row.WriteRune('x')
			} else {
				row.WriteRune(' ')
			}
		}
		rows = append(rows, row.String())
	}
	return strings.Join(rows, "\n")
}

func (s *Shape) Clone() *Shape {
	return ShapeFromString(s.String())
}

func (s *Shape) Rotate90() *Shape {
	var newGrid [][]bool
	for x := 0; x < s.Width(); x++ {
		var column []bool
		for y := 0; y < s.Height(); y++ {
			column = append(column, s.grid[y][x])
		}
		newGrid = append(newGrid, column)
	}
	for y := 0; y < len(newGrid); y++ {
		var row []bool
		for x := len(newGrid[0]) - 1; x >= 0; x-- {
			row = append(row, newGrid[y][x])
		}
		newGrid[y] = row
	}
	s.grid = newGrid
	return s
}

func (s *Shape) Flip() *Shape {
	var newGrid [][]bool
	for y := 0; y < s.Height(); y++ {
		var row []bool
		for x := s.Width() - 1; x >= 0; x-- {
			row = append(row, s.grid[y][x])
		}
		newGrid = append(newGrid, row)
	}
	s.grid = newGrid
	return s
}

func ShapeGroups(shapes []*Shape) [][]*Shape {
	var shapeGroups [][]*Shape
	for _, shape := range shapes {
		var shapeGroup []*Shape
		allPermutations := []*Shape{
			shape,
			shape.Clone().Rotate90(),
			shape.Clone().Rotate90().Rotate90(),
			shape.Clone().Rotate90().Rotate90().Rotate90(),
			shape.Clone().Flip(),
			shape.Clone().Flip().Rotate90(),
			shape.Clone().Flip().Rotate90().Rotate90(),
			shape.Clone().Flip().Rotate90().Rotate90().Rotate90(),
		}
		for _, permutation := range allPermutations {
			s := permutation.String()
			same := false
			for _, shape := range shapeGroup {
				if s == shape.String() {
					same = true
					break
				}
			}
			if same {
				continue
			}
			permutation.ComputePoints()
			shapeGroup = append(shapeGroup, permutation)
		}
		shapeGroups = append(shapeGroups, shapeGroup)
	}
	return shapeGroups
}

type Point struct {
	x int
	y int
}

type Solution struct {
	grid [][]rune
}

func (s *Solution) Width() int {
	return len(s.grid[0])
}

func (s *Solution) Height() int {
	return len(s.grid)
}

func (s *Solution) String() string {
	var rows []string
	for y := 0; y < s.Height(); y++ {
		var row strings.Builder
		for x := 0; x < s.Width(); x++ {
			row.WriteRune(s.grid[y][x])
		}
		rows = append(rows, row.String())
	}
	return strings.Join(rows, "\n")
}

func SolutionFromPoints(points []Point) *Solution {
	var solution Solution
	maxX := 0
	maxY := 0
	for _, point := range points {
		if point.x > maxX {
			maxX = point.x
		}
		if point.y > maxY {
			maxY = point.y
		}
	}
	for y := 0; y <= maxY; y++ {
		var row []rune
		for x := 0; x <= maxX; x++ {
			hasPoint := false
			for _, point := range points {
				if point.x == x && point.y == y {
					hasPoint = true
					break
				}
			}
			value := ' '
			if hasPoint {
				value = 'x'
			}
			row = append(row, value)
		}
		solution.grid = append(solution.grid, row)
	}
	return &solution
}

func SolvePuzzle(puzzle *Shape, shapes []*Shape) {
	// Assumptions - puzzle width & height both bigger than any shape width & height.
	puzzle.ComputePoints()
	shapeGroups := ShapeGroups(shapes)

	// Array of shape group index, shape index, and offset index
	// A shape might have no offsets, in which case the offset list may be nil.
	var possibleFits [][][]Point

	for _, shapeGroup := range shapeGroups {
		var shapeGroupFits [][]Point
		for _, shape := range shapeGroup {
			var shapeFits []Point
			for yOffset := 0; yOffset <= puzzle.Height()-shape.Height(); yOffset++ {
				for xOffset := 0; xOffset <= puzzle.Width()-shape.Width(); xOffset++ {
					fits := true
					for y := 0; y < shape.Height(); y++ {
						for x := 0; x < shape.Width(); x++ {
							if shape.grid[y][x] && !puzzle.grid[y+yOffset][x+xOffset] {
								fits = false
								break
							}
						}
						if !fits {
							break
						}
					}
					if fits {
						shapeFits = append(shapeFits, Point{x: xOffset, y: yOffset})
					}
				}
			}
			shapeGroupFits = append(shapeGroupFits, shapeFits)
		}
		possibleFits = append(possibleFits, shapeGroupFits)
	}

	solution := SolutionFromPoints(puzzle.points)
	// The following loop contains the most important optimization.
	// Prune the offsets based on which shape offsets would make some locations
	// impossible to reach through any of the other shapes' placements.
	// The key insight here is that we can prune entire offsets before recursing.
	// This final change improved the runtime from around 8s to 18ms.
	for shapeGroupIndex, shapeGroup := range possibleFits {
		for shapeIndex, offsets := range shapeGroup {
			var validOffsetIndicies []int
			for offsetIndex, offset := range offsets {
				points := shapeGroups[shapeGroupIndex][shapeIndex].points
				// Demark this offset's placement as `1`.
				for _, point := range points {
					solution.grid[point.y+offset.y][point.x+offset.x] = '1'
				}

				for otherShapeGroupIndex, shapeGroup := range possibleFits {
					if shapeGroupIndex == otherShapeGroupIndex {
						continue
					}
					for shapeIndex, offsets := range shapeGroup {
						for _, offset := range offsets {
							points := shapeGroups[otherShapeGroupIndex][shapeIndex].points
							// Ensure that this other piece's offset will not collide with the `1`s.
							fits := true
							for _, point := range points {
								if solution.grid[point.y+offset.y][point.x+offset.x] == '1'  {
									fits = false
									break
								}
							}
							if !fits {
								continue
							}
							// Demark the reached cells from this piece as `2`.
							for _, point := range points {
								solution.grid[point.y+offset.y][point.x+offset.x] = '2'
							}
						}
					}
				}

				// At this point, we have marked the offset we're looking at as `1`,
				// and all the other reachable places as `2`.
				// Any places still with `x` are unreachable if we use this offset, so we can prune it.
				impossible := false
				for _, point := range puzzle.points {
					if solution.grid[point.y][point.x] == 'x' {
						impossible = true
						break
					}
				}
				if !impossible {
					validOffsetIndicies = append(validOffsetIndicies, offsetIndex)
				}

				// Undo our modifications to the grid.
				for _, point := range puzzle.points {
					solution.grid[point.y][point.x] = 'x'
				}
			}
			var validOffsets []Point
			for _, validOffsetIndex := range validOffsetIndicies {
				validOffsets = append(validOffsets, offsets[validOffsetIndex])
			}
			possibleFits[shapeGroupIndex][shapeIndex] = validOffsets
		}
	}

	// Sort shape groups by complexity - least complex in the front, most complex at the end.
	// This change reduced the runtime from about 11s to 8s.
	// With all other optimizations, the default order runs in about 28ms.
	// Least optimal order runs in 138ms.
	// Most optimal order runs in 18ms.
	var order []int
	for i := range possibleFits {
		order = append(order, i)
	}
	sort.Slice(order, func(i, j int) bool {
		a := 0
		b := 0
		for _, shape := range possibleFits[i] {
			a += len(shape)
		}
		for _, shape := range possibleFits[j] {
			b += len(shape)
		}
		return a < b
	})

	RecursiveSolve(solution, shapeGroups, possibleFits, 0, order)
	fmt.Println(solution)
}

func RecursiveSolve(solution *Solution, shapeGroups [][]*Shape, possibleFits [][][]Point, shapeGroupIndex int, order []int) bool {
	if shapeGroupIndex == len(shapeGroups) {
		// Base case, solved.
		return true
	}

	for shapeIndex, offsets := range possibleFits[order[shapeGroupIndex]] {
		points := shapeGroups[order[shapeGroupIndex]][shapeIndex].points
		for _, offset := range offsets {
			// Check if it fits.
			offsetFits := true
			for _, point := range points {
				if solution.grid[point.y+offset.y][point.x+offset.x] != 'x' {
					// This one doesn't fit.
					offsetFits = false
					break
				}
			}
			if !offsetFits {
				// Try something else.
				continue
			}

			for _, point := range points {
				solution.grid[point.y+offset.y][point.x+offset.x] = rune(shapeGroupIndex + 65)
			}
			// Recurse with this fit.
			if RecursiveSolve(solution, shapeGroups, possibleFits, shapeGroupIndex+1, order) {
				return true
			}
			// Backtrack.
			for _, point := range points {
				solution.grid[point.y+offset.y][point.x+offset.x] = 'x'
			}
		}
	}

	return false
}

func main() {
	SolvePuzzle(
		ShapeFromString(
			"                 xx \n"+
				"               xxxxx\n"+
				"              xx    \n"+
				"             xx     \n"+
				"            xxx     \n"+
				"       xxx xxx      \n"+
				"    xxxxxxxxxx      \n"+
				"   xxxxxxxxxx       \n"+
				"xxxxxxxxxxxx        \n"+
				"     xx   xx        \n"+
				"     xx   xx        "),
		[]*Shape{
			ShapeFromString("xxxxx"),
			ShapeFromString(" xxx\nxx  "),
			ShapeFromString("xx\nx \nxx"),
			ShapeFromString(" x \nxxx\n x "),
			ShapeFromString("xxx\n x \n x "),
			ShapeFromString("xx\nx \nx \nx "),
			ShapeFromString("xxx\nx  \nx  "),
			ShapeFromString("xx \n x \n xx"),
			ShapeFromString("x \nxx\nx \nx "),
			ShapeFromString(" xx\nxx \nx  "),
			ShapeFromString("xx \nxxx"),
			ShapeFromString(" xx\nxx \n x "),
		},
	)
}