← Home
For problem statement at 1000-1999/1400-1499/1420-1429/1428/problemE.txt this is a correct solution, but verifier at 1000-1999/1400-1499/1420-1429/1428/verifierE.go ends with All 100 tests passed can you fix the verifier? package main

import (
	"bufio"
	"container/heap"
	"io"
	"os"
	"strconv"
)

type FastScanner struct {
	data []byte
	idx  int
	n    int
}

func NewFastScanner() *FastScanner {
	data, _ := io.ReadAll(os.Stdin)
	return &FastScanner{data: data, n: len(data)}
}

func (fs *FastScanner) NextInt() int {
	for fs.idx < fs.n && (fs.data[fs.idx] < '0' || fs.data[fs.idx] > '9') {
		fs.idx++
	}
	val := 0
	for fs.idx < fs.n && fs.data[fs.idx] >= '0' && fs.data[fs.idx] <= '9' {
		val = val*10 + int(fs.data[fs.idx]-'0')
		fs.idx++
	}
	return val
}

func cost(a, p int) int64 {
	q := a / p
	r := a % p
	qq := int64(q)
	return int64(p-r)*qq*qq + int64(r)*(qq+1)*(qq+1)
}

func gain(a, p int) int64 {
	if p >= a {
		return 0
	}
	return cost(a, p) - cost(a, p+1)
}

type Node struct {
	a int
	p int
	d int64
}

type MaxHeap []Node

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i].d > h[j].d }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *MaxHeap) Push(x interface{}) {
	*h = append(*h, x.(Node))
}

func (h *MaxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func main() {
	fs := NewFastScanner()
	n := fs.NextInt()
	k := fs.NextInt()

	h := make(MaxHeap, n)
	var total int64

	for i := 0; i < n; i++ {
		a := fs.NextInt()
		total += int64(a) * int64(a)
		h[i] = Node{a: a, p: 1, d: gain(a, 1)}
	}

	heap.Init(&h)

	for i := 0; i < k-n; i++ {
		x := heap.Pop(&h).(Node)
		total -= x.d
		x.p++
		x.d = gain(x.a, x.p)
		heap.Push(&h, x)
	}

	w := bufio.NewWriterSize(os.Stdout, 1<<20)
	w.WriteString(strconv.FormatInt(total, 10))
	w.Flush()
}