E - Apple Baskets on Circle Editorial by cirno3153

\(O(N \log N)\) で解きます。

高橋君の行動を高速にシミュレーションします。 まず、高橋くんがかご \(1\) に戻ってくる回数を考えます。

りんごが \(1\) 個以上のかごの個数を \(C\) とし、りんごが \(0\) 個のかごをいったん除外します。 すると、高橋君は \(\min(A)\) 回の間は \(C\) 個のかごからそれぞれ \(1\) 個ずつりんごを食べてかご \(1\) に戻ることを繰り返します。

従って、 \(C \min(A) \leq K\) の場合、まずりんごが入っているかご全てから \(C\) 個ずつりんごを食べることでかごが一つ減ります。

逆に、 \(C\min(A) > K\) の場合、まずりんごが入っているかご全てから \(\lfloor \frac{K}{C} \rfloor\) 個ずつりんごを食べることで \(K\) が \(N\) 未満になります。

\(K\) が \(N\) 未満になった場合、後は愚直にシミュレーションしても \(O(N)\) です。 そうでない場合、要は周回をしたということなので周回回数だけ管理すれば後で \(O(N)\) で計算できます。

後は \(\min(A)\) を適切に更新する部分ですが、最小のものから順に取り除くので優先度付きキューがあれば求まります。

C++の実装例

#include<iostream>
#include<queue>
#include<vector>
using namespace std;
using ll = long long int;
int main() {
  int N;
  ll K;
  cin >> N >> K;
  vector<ll> A(N);
  priority_queue<ll, vector<ll>, greater<ll>> pq;
  for (int i = 0;i < N;++ i) {
    cin >> A[i];
    pq.push(A[i]);
  }
  ll loop = 0; // この回数だけ周回する
  for (int i = N;;-- i) { // i: まだりんごが入っているかごの数
    ll nextLoop = pq.top() - loop; // 周回回数
    pq.pop();
    if (nextLoop * i > K) nextLoop = K / i;
    K -= nextLoop * i;
    loop += nextLoop;
    if (K < i) break; // 後は微調整で行ける
  }
  for (long a : A) {
    ll ans = max(0LL, a - loop); // まず周回分を除外
    if (a > loop && K --> 0LL) -- ans; // まだK個食べていないなら追加で食べる
    cout << ans << endl;
  }
  return 0;
}

Javaの実装例

import java.util.PriorityQueue;
import java.util.Scanner;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import static java.lang.System.out;
public class Main {
  public static void main(String[] args) {
    try (Scanner sc = new Scanner(System.in)) {
      int N = sc.nextInt();
      long K = sc.nextLong();
      long[] A = IntStream.range(0, N).mapToLong(i -> sc.nextLong()).toArray();
      PriorityQueue<Long> pq = new PriorityQueue<>(IntStream.range(0, N).mapToObj(i -> A[i]).collect(Collectors.toList()));
      long loop = 0; // この回数だけ周回する
      for (int i = N;;-- i) { // i: まだりんごが入っているかごの数
        long nextLoop = pq.poll() - loop; // 周回回数
        if (nextLoop * i > K) nextLoop = K / i;
        K -= nextLoop * i;
        loop += nextLoop;
        if (K < i) break; // 後は微調整で行ける
      }
      for (long a : A) {
        long ans = Math.max(0L, a - loop); // まず周回分を除外
        if (a > loop && K --> 0L) -- ans; // まだK個食べていないなら追加で食べる
        out.println(ans);
      }
    }
  }
}

Kotlinの実装例

import java.util.PriorityQueue
import java.util.Scanner
import kotlin.math.max
fun main() {
  Scanner(System.`in`).use { sc -> 
    val N = sc.nextInt()
    var K = sc.nextLong()
    val A = LongArray(N) {sc.nextLong()}
    val pq = PriorityQueue(A.toList())
    var loop = 0L // この回数だけ周回する
    for (i in N downTo 0) { // i: まだりんごが入っているかごの数
      var nextLoop = pq.poll() - loop // 周回回数
      if (nextLoop * i > K) nextLoop = K / i
      K -= nextLoop * i
      loop += nextLoop
      if (K < i) break // 後は微調整で行ける
    }
    for (a in A) {
      var ans = max(0L, a - loop) // まず周回分を除外
      if (a > loop && K --> 0L) -- ans // まだK個食べていないなら追加で食べる
      println(ans)
    }
  }
}

Pythonの実装例

import copy
import heapq
N, K = map(int, input().split())
A = list(map(int, input().split()))
pq = copy.copy(A)
heapq.heapify(pq)
loop = 0 # この回数だけ周回する
for i in reversed(range(1, N + 1)): # i: まだりんごが入っているかごの数
  nextLoop = heapq.heappop(pq) - loop # 周回回数
  if nextLoop * i > K: nextLoop = K // i
  K -= nextLoop * i
  loop += nextLoop
  if K < i: break # 後は微調整で行ける
for a in A:
  ans = max(0, a - loop) # まず周回分を除外
  if a > loop and K > 0: # まだK個食べていないなら追加で食べる
    ans = ans - 1
    K = K - 1
  print(ans)

posted:
last update: