Official

C - Fair Candy Distribution Editorial by Nyaan


問題で与えられた操作を愚直に行うと、最大で \(K = 10^{18}\) 回の操作が必要になりTLEしてしまうので、適切な考察により計算量を改善する必要があります。

全員に配る回数は全部で何回になるかを考えると、\(K \times t\)\(N\) を超えないような最大の \(t\) が答えになります。この \(t\)

\[K \times t \leq N \iff t \leq \frac{N}{K}\]

より、\(N\)\(K\) で切り捨て除算した値 \(\left \lfloor \frac{N}{K} \right \rfloor\) になります。

以上より全員に配る部分は高速に操作することが出来ました。残りのお菓子の個数は \(K' = N - K \times \left \lfloor \frac{N}{K} \right \rfloor\) 個で、この \(K'\) 個を番号の小さい \(K'\) 人に配ればこの問題を解くことが出来ます。

実装方法はいくつかありますが、ここでは他の問題でも使える \(\mathrm{O}(N \log N)\) のアルゴリズムを説明します。

  1. 新たに配列 order を宣言して、 \(1,2,\ldots,N\) を追加する。
  2. \(i \lt j\) ならば \(\mathrm{a} \lbrack i \rbrack \lt \mathrm{a} \lbrack j \rbrack\) になるように比較関数を定義して、配列 order をソートする。
  3. 配列 order の先頭 \(K'\) 項の添え字に対応する要素が求める要素である。

以上のアルゴリズムにより、 \(a_i\) の小さい \(K'\) 人が誰かを高速に求めることが出来ました。

これは余談ですが、 Quick Select や Median of Medians などの選択アルゴリズムを使用すると \(\mathrm{O}(N)\) での計算が可能になります。興味がある方は調べてみると勉強になると思います。

C++ と Python による実装例は次の通りです。 Python の場合はナイーブな実装と定数倍の良い実装の二つを載せています。(解説では 1-indexed で説明した部分が 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
# ナイーブな実装
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)
# 定数倍の良い実装
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: