diff --git a/cmd/meanstoanend/main.go b/cmd/meanstoanend/main.go new file mode 100644 index 0000000..e49d627 --- /dev/null +++ b/cmd/meanstoanend/main.go @@ -0,0 +1,60 @@ +package main + +import ( + "encoding/binary" + "io" + "log" + "net" + + "protohackers/pkg/conn" +) + +func main() { + err := conn.StartSimple(meansToAnEnd) + if err != nil { + log.Fatalln(err) + } +} + +func meansToAnEnd(c net.Conn) { + defer c.Close() + var ledger = make(map[int32]int32) + buf := make([]byte, 9) + + for { + if _, err := io.ReadFull(c, buf); err != nil { + return + } + if _, err := c.Write(process(buf, ledger)); err != nil { + return + } + } +} + +func process(buf []byte, ledger map[int32]int32) []byte { + i1, i2 := int32(binary.BigEndian.Uint32(buf[1:5])), int32(binary.BigEndian.Uint32(buf[5:])) + + switch buf[0] { + case 'I': + ledger[i1] = i2 + case 'Q': + var c int64 + var t int64 + for k, v := range ledger { + if i1 <= k && k <= i2 { + t += int64(v) + c++ + } + } + if c == 0 { + c = 1 /* Can't divide by zero! */ + } + mean := t / c + r := make([]byte, 4) + binary.BigEndian.PutUint32(r, uint32(mean)) + return r + default: + return nil + } + return nil +} diff --git a/cmd/primetime/main.go b/cmd/primetime/main.go index 5921b0e..17e85fe 100644 --- a/cmd/primetime/main.go +++ b/cmd/primetime/main.go @@ -1,7 +1,6 @@ package main import ( - "bufio" "encoding/json" "log" "math/big" @@ -29,13 +28,7 @@ func main() { } func primetime(c net.Conn) { - defer c.Close() - sc := bufio.NewScanner(c) - for sc.Scan() { - response := process(sc.Bytes()) - c.Write(response) - } - if err := sc.Err(); err != nil { + if err := conn.ReadAndRespond(c, process); err != nil { log.Fatalln(err) } } diff --git a/pkg/conn/conn.go b/pkg/conn/conn.go index 9d266e9..7f03be8 100644 --- a/pkg/conn/conn.go +++ b/pkg/conn/conn.go @@ -1,6 +1,7 @@ package conn import ( + "bufio" "fmt" "net" ) @@ -30,3 +31,19 @@ func StartSimple(handler func(conn net.Conn)) error { } } } + +func ReadAndRespond(c net.Conn, process func([]byte) []byte) error { + defer c.Close() + sc := bufio.NewScanner(c) + for sc.Scan() { + response := process(sc.Bytes()) + _, err := c.Write(response) + if err != nil { + return err + } + } + if err := sc.Err(); err != nil { + return err + } + return nil +}