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):
  answer[order[i].id] += 1
for x in answer:
  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
for x in answer:
  print(x)

posted:
last update: