protohackers/cmd/meanstoanend/main.go

package main

import (
	"encoding/binary"
	"io"
	"log"
	"net"

	"protohackers/pkg/conn"
)

func main() {
	err := conn.StartSimpleTCP(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
}