AtCoder Beginner Contest 216E

ABC216E

同じアトラクションに何回乗るのかを高速に求められれば十分．
Multiset としての和 \(B := \cup_{i \in N} \{1, \cdots, a_{i}\}\) を取る． \(B\) の元から，大きいほうから \(K\) 個まで取ってくるのが最善．
\(m \in \mathbb{N}\) に対して， \(Bool \ni f(m) := \) (\(B\) に含まれる \(m\) 以上の値は \(K\) 個以下) とする． \(C\) を十分大きな定数とする．\(C = max_{i \in N} (a_{i}) + 1\) など． 以下，最適なアトラクションの乗り方を高速に求める．

Case : \(\vert B \vert \leq K\)
このときは，\(B\) からすべて取ればよい． ちなみに \(f(0) = true, f(C) = true\).

Case : otherwise
\(f(0) = false, f(C) = true\). Binary search により，\(f(m) = true\) となる minimum な \(m\) が求まるので， それを \(m_{0}\) とおく．
各 \(i \in N\) に対して， if \(a_{i} \geq m_{0}\)のとき \(a_{i}, a_{i} -1, \cdots, m_{0}\) までを \(ans\) に加算して， \(K\) から \(a_{i} + 1 - m_{0} \) を引く．
\(i\) 毎の操作が終わった後，\(K > 0\) の場合は， 楽しさ \(m_{0} -1\) のアトラクションを \(K\) 回乗れる．

使っている記号，マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

int main() {

  ll n, k ;

  cin >> n >> k;

  vll a(n); rep(i,n) { cin >> a[i]; }

  ll ok, ng;

  ok = 1e18, ng = 0;

  while(abs(ok-ng) > 1){

    ll md = (ok + ng)/2;

    ll c = 0;

    rep(i,n){

      c += max(a[i]-md+1, 0LL);

      // c += a[i] - min(md, a[i]) + 1;

    }

    if(c <= k) ok = md;

    else ng = md;

  }

 

  ll ans = 0;

  rep(i,n){

    ll l = a[i];

    ll r = ok-1;

    // [l,r)

    auto sum = [&](ll l, ll r, ll d) -> ll {

      ll len = (r-l)/d;

      return len*(l+r-d)/2;

    };

    if(l > r){

      ans += sum(l,r,-1);

      k -= abs(l-r);

    }

  }

  assert(k >= 0);

  ans += k * (ok-1);

  k = 0;

  cout << ans << endl;

 

  return 0;

}