Official

G - Redistribution of Piles Editorial by Nyaan


2 種類の操作をそれぞれ操作 \(A\), 操作 \(B\) と呼びます。 重要な観察として次の事実があります。

  • 「操作 \(B\) の直後に操作 \(A\) を行わない」という条件を加えても答えは変わらない。

なぜかというと、操作 \(B\) の直後は全ての袋にボールが 1 個ずつ入っているので操作 \(A\) を行うと元の状態に戻ってしまうからです。よってこのような操作は無意味であると分かります。
同様にして以下の事実も確認できます。

  • 「操作 \(A\) で全ての袋からボールを取り出した直後に, 操作 \(B\) を行わない」という条件を加えても答えは変わらない。

また、逆に上記の 2 つのルールを加えた場合に、 1 つの操作列と 1 個の数列の間に全単射を取れることが確認できます。よってこの問題は条件を満たす操作列の数え上げに言い換えられます。

操作列を整数 \(x, y\) を用いて 「\(A\)\(x\) 回, \(B\)\(y\) 回」 という形で表すとします。\(x\) を固定したときの \(y\) の範囲と通り数を考えると以下のようになります。

  • \(x \leq \min \lbrace A \rbrace\) の場合:\(y = 0\) (\(1\) 通り) (\(B\) を行うと 2 番目の条件に反する)
  • \(\min \lbrace A \rbrace \lt x\) の場合 : \(s\) を袋に入っているボールの個数として, \(0 \leq y \leq \lfloor s/N \rfloor\) (\( \lfloor s/N \rfloor + 1\) 通り)

よって \(x\) を全探索すると \(\mathrm{O}(\max \lbrace A \rbrace)\) でこの問題を解けますが, 計算量が大きすぎて TLE します。
式を整理すると \(\sum_{L \leq n \leq R} \lfloor (ax+b)/m \rfloor\) 型の式が \(\mathrm{O}(N)\) 個ある形に変形できるため、floor sum と呼ばれるアルゴリズムを適用できる形に帰着できます。 (floor sum は AtCoder Library を使えば容易に計算できます。参考:ACL のリファレンス )
floor sum を利用するとこの問題を (\(M = \max \lbrace A \rbrace\)として) \(\mathrm{O}(N \log M)\) で計算出来て十分高速です。

  • 実装例(C++)
#include <bits/stdc++.h>
using namespace std;

#include "atcoder/math.hpp"
#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;
using ll = long long;

int main() {
  ll N;
  cin >> N;
  vector<ll> A(N);
  for(auto&x : A) cin >> x;
  sort(begin(A), end(A));
  ll s = A[0] * N;
  mint ans = A[0] + 1;
  for (int i = 1; i < N; i++) {
    ll L = A[i - 1];
    ll R = A[i];
    ll a = N - i;
    ll b = s - L * (N - i);
    ll m = N;
    ll cur = 0;
    cur += atcoder::floor_sum(R + 1, m, a, b);
    cur -= atcoder::floor_sum(L + 1, m, a, b);
    cur += R - L;
    ans += cur;
    s += (R - L) * (N - i);
  }
  cout << ans.val() << "\n";
}

posted:
last update: