← Home
For problem statement at 0-999/800-899/800-809/802/problemE.txt this is a correct solution, but verifier at 0-999/800-899/800-809/802/verifierE.go ends with All 102 tests passed. can you fix the verifier? package main

import (
	"bytes"
	"bufio"
	"io"
	"math"
	"os"
	"strconv"
)

var factCache = map[int]float64{
	0: 0,
	1: 0,
}

func lnFact(x int) float64 {
	if v, ok := factCache[x]; ok {
		return v
	}
	v, _ := math.Lgamma(float64(x + 1))
	factCache[x] = v
	return v
}

func parseIntsLine(line []byte) []int {
	line = bytes.TrimSpace(line)
	if len(line) == 0 {
		return nil
	}
	res := make([]int, 0, 16)
	sign := 1
	num := 0
	in := false
	for _, c := range line {
		if c == '-' {
			sign = -1
		} else if c >= '0' && c <= '9' {
			num = num*10 + int(c-'0')
			in = true
		} else if in {
			res = append(res, sign*num)
			sign = 1
			num = 0
			in = false
		}
	}
	if in {
		res = append(res, sign*num)
	}
	return res
}

func buildVillages(lines [][]int) [][]int {
	if len(lines) == 0 {
		return nil
	}

	if len(lines[0]) == 1 {
		t := lines[0][0]
		rem := lines[1:]

		if len(rem) >= 1 && len(rem[0]) == 1 {
			m := rem[0][0]
			if len(rem[1:]) == t {
				exact := true
				pref := true
				for _, line := range rem[1:] {
					if len(line) != m {
						exact = false
					}
					if len(line) != m+1 || line[0] != m {
						pref = false
					}
				}
				if exact || pref {
					v := make([][]int, t)
					for i, line := range rem[1:] {
						if pref {
							v[i] = line[1:]
						} else {
							v[i] = line
						}
					}
					return v
				}
			}
		}

		if len(rem) == t {
			strip := true
			for _, line := range rem {
				if len(line) == 0 || line[0] != len(line)-1 {
					strip = false
					break
				}
			}
			v := make([][]int, t)
			for i, line := range rem {
				if strip {
					v[i] = line[1:]
				} else {
					v[i] = line
				}
			}
			return v
		}

		if len(rem) == 2*t {
			ok := true
			for i := 0; i < t; i++ {
				if len(rem[2*i]) != 1 || rem[2*i][0] != len(rem[2*i+1]) {
					ok = false
					break
				}
			}
			if ok {
				v := make([][]int, t)
				for i := 0; i < t; i++ {
					v[i] = rem[2*i+1]
				}
				return v
			}
		}

		flat := make([]int, 0)
		for _, line := range rem {
			flat = append(flat, line...)
		}
		if t > 0 && len(flat)%t == 0 {
			m := len(flat) / t
			v := make([][]int, t)
			idx := 0
			for i := 0; i < t; i++ {
				v[i] = flat[idx : idx+m]
				idx += m
			}
			return v
		}
	}

	if len(lines[0]) == 2 {
		a, b := lines[0][0], lines[0][1]
		rem := lines[1:]

		try := func(t, m int) [][]int {
			if t < 0 || m < 0 || len(rem) != t {
				return nil
			}
			exact := true
			pref := true
			for _, line := range rem {
				if len(line) != m {
					exact = false
				}
				if len(line) != m+1 || line[0] != m {
					pref = false
				}
			}
			if !exact && !pref {
				return nil
			}
			v := make([][]int, t)
			for i, line := range rem {
				if pref {
					v[i] = line[1:]
				} else {
					v[i] = line
				}
			}
			return v
		}

		if v := try(a, b); v != nil {
			return v
		}
		if v := try(b, a); v != nil {
			return v
		}
	}

	return lines
}

func estimate(xs []int) int {
	n := len(xs)
	if n == 0 {
		return 0
	}

	sum := 0.0
	maxVal := 0
	for _, x := range xs {
		sum += float64(x)
		if x > maxVal {
			maxVal = x
		}
	}
	mean := sum / float64(n)

	lp := 0.0
	if mean == 0 {
		for _, x := range xs {
			if x != 0 {
				lp = math.Inf(-1)
				break
			}
		}
	} else {
		logMean := math.Log(mean)
		for _, x := range xs {
			lp += float64(x)*logMean - mean - lnFact(x)
		}
	}

	minP := (maxVal + 1) / 2
	lu := -float64(n) * math.Log(float64(2*minP+1))

	if lu > lp {
		maxCorr := float64(maxVal) * (float64(n) + 1) / (2 * float64(n))
		w := float64(n+2) / float64(n+5)
		est := w*maxCorr + (1-w)*mean
		ans := int(math.Round(est))
		if ans < minP {
			ans = minP
		}
		if ans < 0 {
			ans = 0
		}
		return ans
	}

	ans := int(math.Round(mean))
	if ans < 0 {
		ans = 0
	}
	return ans
}

func main() {
	data, _ := io.ReadAll(os.Stdin)
	rawLines := bytes.Split(data, []byte{'\n'})
	lines := make([][]int, 0, len(rawLines))
	for _, line := range rawLines {
		ints := parseIntsLine(line)
		if len(ints) > 0 {
			lines = append(lines, ints)
		}
	}

	villages := buildVillages(lines)

	out := bufio.NewWriter(os.Stdout)
	for _, v := range villages {
		out.WriteString(strconv.Itoa(estimate(v)))
		out.WriteByte('\n')
	}
	out.Flush()
}