Solve day 6

day6/example.txt

@@ -0,0 +1,4 @@
+123 328  51 64 
+ 45 64  387 23 
+  6 98  215 314
+*   +   *   +  

day6/main.go

@@ -0,0 +1,217 @@
+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
+			}
+		}
+	}
+}