218 lines
3.9 KiB
Go
218 lines
3.9 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"iter"
|
|
"log"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
func main() {
|
|
f := os.Stdin
|
|
if len(os.Args) > 1 && os.Args[1] != "-" {
|
|
var err error
|
|
f, err = os.Open(os.Args[1])
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
defer f.Close()
|
|
}
|
|
|
|
lines, err := readLines(f)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
problems1, err := parseProblemsPart1(lines)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
problems2, err := parseProblemsPart2(lines)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
var result1, result2 int
|
|
for _, p := range problems1 {
|
|
result1 += p.Solve()
|
|
}
|
|
for _, p := range problems2 {
|
|
result2 += p.Solve()
|
|
}
|
|
fmt.Println("part one:", result1)
|
|
fmt.Println("part two:", result2)
|
|
}
|
|
|
|
type Action int
|
|
|
|
const (
|
|
ActionAdd Action = iota
|
|
ActionMultiply
|
|
)
|
|
|
|
type problem struct {
|
|
Values []int
|
|
Action
|
|
}
|
|
|
|
func (p problem) Solve() int {
|
|
if len(p.Values) == 0 {
|
|
return 0
|
|
}
|
|
answer := p.Values[0]
|
|
for _, v := range p.Values[1:] {
|
|
switch p.Action {
|
|
case ActionAdd:
|
|
answer += v
|
|
case ActionMultiply:
|
|
answer *= v
|
|
default:
|
|
panic(fmt.Sprintf("unsupported action %d", p.Action))
|
|
}
|
|
}
|
|
return answer
|
|
}
|
|
|
|
func readLines(r io.Reader) ([]string, error) {
|
|
var lines []string
|
|
sc := bufio.NewScanner(r)
|
|
for sc.Scan() {
|
|
lines = append(lines, sc.Text())
|
|
}
|
|
return lines, sc.Err()
|
|
}
|
|
|
|
func parseFields(lines []string) [][]string {
|
|
cutAll := func(lines []string, sep string) (before, after []string, ok bool) {
|
|
var found bool
|
|
var i int
|
|
for _, line := range lines {
|
|
j := strings.Index(line, sep)
|
|
if j < 0 {
|
|
j = len(line)
|
|
} else {
|
|
found = true
|
|
}
|
|
i = max(i, j)
|
|
}
|
|
if !found {
|
|
return lines, nil, false
|
|
}
|
|
|
|
before = make([]string, 0, len(lines))
|
|
after = make([]string, 0, len(lines))
|
|
for _, line := range lines {
|
|
before = append(before, line[:min(i, len(line))])
|
|
after = append(after, line[min(i+len(sep), len(line)):])
|
|
}
|
|
return before, after, true
|
|
}
|
|
|
|
fields := make([][]string, len(lines))
|
|
for {
|
|
before, after, ok := cutAll(lines, " ")
|
|
if !ok {
|
|
break
|
|
}
|
|
for i, v := range before {
|
|
fields[i] = append(fields[i], v)
|
|
}
|
|
lines = after
|
|
}
|
|
return fields
|
|
}
|
|
|
|
func parseProblemsPart1(lines []string) ([]problem, error) {
|
|
if len(lines) == 0 {
|
|
return nil, nil
|
|
}
|
|
var data [][]string
|
|
for _, line := range lines {
|
|
data = append(data, strings.Fields(line))
|
|
}
|
|
|
|
problems := make([]problem, len(data[0]))
|
|
for i, row := range data {
|
|
if len(row) != len(problems) {
|
|
return nil, fmt.Errorf("line %d: unexpected number of fields", i+2)
|
|
}
|
|
|
|
last := i == len(data)-1
|
|
for j, v := range row {
|
|
p := &problems[j]
|
|
if last {
|
|
switch v {
|
|
case "+":
|
|
p.Action = ActionAdd
|
|
case "*":
|
|
p.Action = ActionMultiply
|
|
default:
|
|
return nil, fmt.Errorf("unexpected action %q", v)
|
|
}
|
|
} else {
|
|
n, err := strconv.Atoi(v)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
p.Values = append(p.Values, n)
|
|
}
|
|
}
|
|
}
|
|
return problems, nil
|
|
}
|
|
|
|
func parseProblemsPart2(lines []string) ([]problem, error) {
|
|
fields := parseFields(lines)
|
|
var problems []problem
|
|
|
|
for col := range columnsSeq(fields) {
|
|
var p problem
|
|
switch s := strings.TrimSpace(col[len(col)-1]); s {
|
|
case "+":
|
|
p.Action = ActionAdd
|
|
case "*":
|
|
p.Action = ActionMultiply
|
|
default:
|
|
return nil, fmt.Errorf("unexpected action %q", s)
|
|
}
|
|
|
|
var colBytes [][]byte
|
|
for _, v := range col[:len(col)-1] {
|
|
colBytes = append(colBytes, []byte(v))
|
|
}
|
|
for v := range columnsSeq(colBytes) {
|
|
s := string(bytes.TrimSpace(v))
|
|
n, err := strconv.Atoi(s)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
p.Values = append(p.Values, n)
|
|
}
|
|
|
|
problems = append(problems, p)
|
|
}
|
|
return problems, nil
|
|
}
|
|
|
|
func columnsSeq[T any](data [][]T) iter.Seq[[]T] {
|
|
return func(yield func([]T) bool) {
|
|
if len(data) < 1 {
|
|
return
|
|
}
|
|
for x := range len(data[0]) {
|
|
var col []T
|
|
for y := range data {
|
|
col = append(col, data[y][x])
|
|
}
|
|
if !yield(col) {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|