Contest Duration: - (local time) (100 minutes) Back to Home

C - Fair Candy Distribution Editorial by en_translator

A naive implementation of operations given in the problem statement requires at most $$K = 10^{18}$$ times of loops, which leads to TLE, so the computation time should be improved with an appropriate consideration.

The number of distributions to everyone is the maximum $$t$$ such that $$K \times t$$ does not exceed $$N$$. Since

$K \times t \leq N \iff t \leq \frac{N}{K},$

the value $$t$$ is equal to the quotient of $$N$$ divided by $$K$$, rounded down, namely $$\left \lfloor \frac{N}{K} \right \rfloor$$.

Therefore, we could compute the distributions to everyone fast enough. The number of remaining candies is $$K' = N - K \times \left \lfloor \frac{N}{K} \right \rfloor$$; now we want to distribute these $$K'$$ candies to $$M$$ citizens with the smallest IDs.

There are some ways of implementation. Here we introduce an $$\mathrm{O}(N \log N)$$ algorithm that can be adoptable for other problems.

1. Prepare an array order and add $$1,2,\ldots,N$$ to it.
2. Sort the array order by a comparison function defined as $$i \lt j$$ implies $$\mathrm{a} \lbrack i \rbrack \lt \mathrm{a} \lbrack j \rbrack$$.
3. The desired elements are those corresponding to the indices of the first $$K'$$ terms of the array order.

With the algorithm above, we have found who are the $$K'$$ citizens with the smallest $$a_i$$.

Interestingly enough, selection algorithms like “quickselect” or “median of medians” enable $$\mathrm{O}(N)$$ computation time. If you are interested, learning it may help you.

The sample code by C++ and Python follows. For Python, we have shown two implementations; one is a naive one and the other is one with a good constant factor. (Note that the editorial is described 1-indexed while the codes are implemented 0-indexed.)

• C++
#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>
using namespace std;

int main() {
int N;
long long K;
cin >> N >> K;

vector<int> a(N);
for (auto& x : a) cin >> x;

vector<int> order(N);
iota(begin(order), end(order), 0);
sort(begin(order), end(order), [&](int i, int j) { return a[i] < a[j]; });

vector<long long> answer(N, K / N);
for(int i = 0; i < K % N; i++) answer[order[i]]++;

for(auto&x : answer) cout << x << "\n";
}
• Python
# A naive implementation
N, K = map(int,input().split())
a = list(map(int,input().split()))

class C:
def __init__(self, i):
self.id = i
def __lt__(self, other):
return a[self.id] < a[other.id]

order = [C(i) for i in range(N)]
order.sort()

answer = [K // N for i in range(N)]
for i in range(K % N):
print(x)
# An implementation with a good constant factor
N, K = map(int,input().split())
a = list(map(int,input().split()))

order = [(a[i] << 32) + i for i in range(N)]
order.sort()

answer = [K // N for i in range(N)]
for i in range(K % N):
answer[order[i] & ((1 << 32) - 1)] += 1