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)\) のアルゴリズムを説明します。
- 新たに配列
order
を宣言して、 \(1,2,\ldots,N\) を追加する。 - \(i \lt j\) ならば \(\mathrm{a} \lbrack i \rbrack \lt \mathrm{a} \lbrack j \rbrack\) になるように比較関数を定義して、配列
order
をソートする。 - 配列
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: