Carrots for Rabbits (Codeforces 1428E)
Tags: Greedy, binary search

There are some rabbits in Singapore Zoo. To feed them, Zookeeper bought $$$n$$$ carrots with lengths $$$a_1, a_2, a_3, \ldots, a_n$$$. However, rabbits are very fertile and multiply very quickly. Zookeeper now has $$$k$$$ rabbits and does not have enough carrots to feed all of them. To solve this problem, Zookeeper decided to cut the carrots into $$$k$$$ pieces. For some reason, all resulting carrot lengths must be positive integers.Big carrots are very difficult for rabbits to handle and eat, so the time needed to eat a carrot of size $$$x$$$ is $$$x^2$$$.Help Zookeeper split his carrots while minimizing the sum of time taken for rabbits to eat the carrots.

Input: The first line contains two integers $$$n$$$ and $$$k$$$ $$$(1 \leq n \leq k \leq 10^5)$$$: the initial number of carrots and the number of rabbits.The next line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ $$$(1 \leq a_i \leq 10^6)$$$: lengths of carrots.It is guaranteed that the sum of $$$a_i$$$ is at least $$$k$$$.

Output: Output one integer: the minimum sum of time taken for rabbits to eat carrots.

Note: For the first test, the optimal sizes of carrots are $$$\{1,1,1,2,2,2\}$$$. The time taken is $$$1^2+1^2+1^2+2^2+2^2+2^2=15$$$For the second test, the optimal sizes of carrots are $$$\{4,5,5,5\}$$$. The time taken is $$$4^2+5^2+5^2+5^2=91$$$.

Sample Inputs:
3 6
5 3 1

1 4
19

Sample Outputs:
15

91

Solution (Python):
import sys
from heapq import heapify, heapreplace
(n, k, *aaa) = map(int, sys.stdin.buffer.read().split())
ans = 0
hq = []
for a in aaa:
    ans += a * a
for a in aaa:
    p = 1
    (d, m) = divmod(a, p + 1)
    cost = d * d * (p + 1 - m) + (d + 1) * (d + 1) * m
    hq.append((cost - a * a, a * a, a, p))
heapify(hq)
for _ in range(k - n):
    (dec, cur, a, p) = hq[0]
    ans += dec
    cur += dec
    p += 1
    (d, m) = divmod(a, p + 1)
    cost = d * d * (p + 1 - m) + (d + 1) * (d + 1) * m
    heapreplace(hq, (cost - cur, cur, a, p))
print(ans)

Solution (Python):
import sys
from heapq import heapify, heapreplace
(n, k, *aaa) = map(int, sys.stdin.buffer.read().split())
ans = 0
for a in aaa:
    ans += a * a
hq = []
for a in aaa:
    p = 1
    (d, m) = divmod(a, p + 1)
    cost = d * d * (p + 1 - m) + (d + 1) * (d + 1) * m
    hq.append((cost - a * a, a * a, a, p))
heapify(hq)
for _ in range(k - n):
    (dec, cur, a, p) = hq[0]
    ans += dec
    cur += dec
    p += 1
    (d, m) = divmod(a, p + 1)
    cost = d * d * (p + 1 - m) + (d + 1) * (d + 1) * m
    heapreplace(hq, (cost - cur, cur, a, p))
print(ans)

Solution (Python):
import heapq

def value(l, parts):
    len1 = l // parts
    len2 = len1 + 1
    cnt2 = l % parts
    cnt1 = parts - cnt2
    return cnt1 * len1 * len1 + cnt2 * len2 * len2

def solve(n, k, a):
    pq = []
    tot = 0
    for i in range(n):
        tot += a[i] * a[i]
        heapq.heappush(pq, (-value(a[i], 1) + value(a[i], 2), a[i], 2))
    for i in range(k - n):
        temp = heapq.heappop(pq)
        tot += temp[0]
        (a, b) = (temp[1], temp[2])
        heapq.heappush(pq, (-value(a, b) + value(a, b + 1), a, b + 1))
    return tot
(n, k) = map(int, input().split())
a = list(map(int, input().split()))
print(solve(n, k, a))