advent of code 2020 problem 11 solution

p11.go

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"reflect"
	"strings"
)

type seatState int

type grid struct {
	seats         []seatState
	columns, rows int
}

const (
	empty seatState = iota
	occupied
	floor
)

func seatRule(seat seatState, adjacents []seatState, maxOccupieds int) seatState {
	occupieds := 0
	for _, s := range adjacents {
		if s == occupied {
			occupieds++
		}
	}
	if seat == empty && occupieds == 0 {
		return occupied
	}
	if seat == occupied && occupieds >= maxOccupieds {
		return empty
	}
	return seat
}

func seatRule1(seat seatState, adjacents []seatState) seatState {
	return seatRule(seat, adjacents, 4)
}

func seatRule2(seat seatState, adjacents []seatState) seatState {
	return seatRule(seat, adjacents, 5)
}

func parse(s string) grid {
	var g grid
	for i, rune := range s {
		switch rune {
		case 'L':
			g.seats = append(g.seats, empty)
		case '.':
			g.seats = append(g.seats, floor)
		case '#':
			g.seats = append(g.seats, occupied)
		case '\n':
			if g.columns == 0 {
				g.columns = i
			}
		}
	}
	g.rows = len(g.seats) / g.columns
	return g
}

func (g grid) seat(x, y int) seatState {
	return g.seats[x*g.columns+y]
}

func (g grid) setSeat(x, y int, s seatState) {
	g.seats[x*g.columns+y] = s
}

func clamp(x, min, max int) int {
	if x < min {
		return min
	}
	if x > max {
		return max
	}
	return x
}

func adjacents(g *grid, x, y int) []seatState {
	var adjacents []seatState
	for i := clamp(-1+x, 0, g.rows-1); i <= clamp(1+x, 0, g.rows-1); i++ {
		for j := clamp(-1+y, 0, g.columns-1); j <= clamp(1+y, 0, g.columns-1); j++ {
			if i == x && j == y {
				continue
			}
			adjacents = append(adjacents, g.seat(i, j))
		}
	}
	return adjacents
}

func directions(g *grid, x, y int) []seatState {
	var adjacents []seatState
	for i := -1; i <= 1; i++ {
		for j := -1; j <= 1; j++ {
			if (i == 0 && j == 0) || (x+i < 0 || x+i >= g.rows || y+j < 0 || y+j >= g.columns) {
				continue
			}
			_x := x + i
			_y := y + j
			for _x >= 0 && _x < g.rows && _y >= 0 && _y < g.columns && g.seat(_x, _y) == floor {
				_x += i
				_y += j
			}
			if _x >= 0 && _x < g.rows && _y >= 0 && _y < g.columns {
				adjacents = append(adjacents, g.seat(_x, _y))
			}
		}
	}
	return adjacents
}

type neighbours func(g *grid, x, y int) []seatState

type rule func(seat seatState, seats []seatState) seatState

func step(g *grid, rule rule, neighbours neighbours) {
	var g2 grid
	g2.seats = make([]seatState, len(g.seats))
	g2.rows = g.rows
	g2.columns = g.columns
	copy(g2.seats, g.seats)
	for i := 0; i < g.rows; i++ {
		for j := 0; j < g.columns; j++ {
			newSeat := rule(g2.seat(i, j), neighbours(&g2, i, j))
			g.setSeat(i, j, newSeat)
		}
	}
}

func (g grid) String() string {
	var b strings.Builder
	for i := 0; i < g.rows; i++ {
		for j := 0; j < g.columns; j++ {
			var seat string
			switch g.seat(i, j) {
			case empty:
				seat = "L"
			case occupied:
				seat = "#"
			case floor:
				seat = "."
			}
			fmt.Fprintf(&b, "%s", seat)
		}
		fmt.Fprintln(&b)
	}
	return b.String()
}

func loop(g *grid, rule rule, neighbours neighbours) {
	for {
		before := make([]seatState, len(g.seats))
		copy(before, g.seats)
		step(g, rule, neighbours)
		if reflect.DeepEqual(g.seats, before) {
			break
		}
	}
}

func part1(input string) *grid {
	g := parse(input)
	loop(&g, seatRule1, adjacents)
	return &g
}

func part2(input string) *grid {
	g := parse(input)
	loop(&g, seatRule2, directions)
	return &g
}

func main() {
	input, err := ioutil.ReadFile("i11")
	if err != nil {
		log.Fatal(err)
	}
	g := part2(string(input))
	occupieds := 0
	for _, seat := range g.seats {
		if seat == occupied {
			occupieds++
		}
	}
	fmt.Println(occupieds)
}