AGC027 B問題 : Garbage Collector - 忘れても大丈夫

問題. Garbage Collector

数直線上に $n$ 個のゴミが左から順番に $0 < x_1 < \cdots < x_n$ の場所にある．ゴミ箱は原点にあり，すべてのゴミを捨てたい．各動作についてコストがかかり，1つのゴミを拾うのに $X$ ，いくつかのゴミを捨てるのに $X$ ，ゴミを $k$ 個持っているときに1移動するのに $(k + 1)^2$ ．ただし，拾ったゴミをゴミ箱以外に置くことはできない．すべてのゴミを捨てるのにかかる最小のコストを求めよ．

制約： $1 \le n \le 2 \times 10^5$ ，　 $0 < x_1 < \cdots < x_n \le 10^9$ ，　 $1 \le X \le 10^9$

解法. 上手くやると調和級数が見えてくる

解説 [1] を参考にした．解説が分かりやすいのでここでは雑にまとめる．

動作が3つ（拾う，捨てる，移動する）ある．拾うはどのゴミのまとめ方でも，各ゴミに対して1回しか行われないので，どの最終のコストに対しても $N \times X$ だけ加算される．したがって，どのようにゴミを拾って移動するのかの戦略を考えることが重要である．
まず，拾い方であるが $t$ 個のゴミ $x_{i_1} < x_{i_2} < \cdots < x_{i_t}$ があるときに，各ゴミに対して行きに拾うか，帰りに拾うかの2通りが考えられる．しかし，移動コストは持っているゴミの量で増加するので明らかに帰りに拾う方が最終コストが悪くならない．よって，原点から出発して， $x_{i_t}, x_{i_{t - 1}}, \cdots, x_{i_1}$ の順番にゴミを拾いながら帰るのが最適な戦略である．
次に，この戦略でかかるコストを考える．まず，捨てるのは1回だけ行われるのでコスト $X$ である．次に，移動コストを考える．まず，原点から $x_{i_t}$ に移動するときのコストは $x_{i_t}$ であり， $x_{i_t}$ から $x_{i_{t - 1}}$ に $x_{i_t}$ を持ちながら移動するのでコストは $2^2 (x_{i_t} - x_{i_{t - 1}})$ である．同様に考えると移動コストは，
　 $x_{i_t} + 2^2 (x_{i_t} - x_{i_{t - 1}}) + 2^3 (x_{i_{t -2}} - x_{i_{t - 3}}) + \cdots + t ^2 (x_{i_2} - x_{i_1}) + (t + 1)^2 x_{i_1}$
　　 $= 5 x_{i_t} + 5 x_{i_{t - 1}} + 7 x_{i_{t - 2}} + 9 x_{i_{t - 3}} + \cdots + (2 (t - j + 1) + 1) x_{i_j} + \cdots + (2t + 1)x_{i_1}$
となる．すなわち，後ろから何番目に拾うかで距離の係数が大きくなる．この性質に着目すると戦略が見えてくる（はず）．ここで，1回に捨てる $t$ 個のゴミは必ずしも連続しているとは限らないことに注意である（反例が $N = 3$ でも作れる）．
$k$ 回に分けてゴミを回収することを考える．このとき， $k$ 回捨てるので捨てるコストは $k \times X$ である．ここで考える必要があるのは， $n$ 個のゴミを何回目に回収するかを決めることである．すなわち， $n$ 個のゴミを $k$ 個の集合に分割する方法でコストが最小となるものを考える必要がある．分割をひとつ決めると上のようにして各分割の移動コストが求まる．この方法を愚直に計算すると $O(n 2^n)$ となるので分割の方法を上手く考える必要があるが，ここで先程の性質を利用すると高速に計算できる．まず，分割の異なる集合に属するゴミ $x_i < x_{i+1}$ で， $x_i$ は後ろから $t$ 番目， $x_{i+1}$ は後ろから $t'$ 番目に位置するものを考える．ただし， $t \ge t'$ とする．分割のコストの各ゴミの位置の項はちょうど1回しか表れず，その係数はそのゴミが属する集合で後ろから何番目に位置するかで決まるので， $x_i$ と $x_{i+1}$ の係数はそれぞれ $(2t + 1)$ と $(2t' + 1)$ である．ただし，一番後ろのゴミの係数は 5 とする．ここで，分割のコストのこれらのゴミの項は $t x_i + t' x_{i+1}$ となる．ここで， $x_i$ と $x_{i+1}$ の属する集合を互いに入れ替えると，他のゴミの項の係数は同じで， $t' x_{i} + t x_{i + 1}$ だけが異なり，交換後の分割のコストは悪くならない．よって，後ろのゴミから順番に $1, 2, \ldots, k, 1, 2, \ldots, k, 1, \cdots$ 番目の集合に属するように分割を行うのが最適な戦略となる．回収する順番が同じゴミは連続しており係数が同じなのでまとめることができる．
例えば， $k = 3$ で $n = 11$ のとき，最適な分割のコストは，
　 $3 \times X + 5 \times (x_{11}+ x_{10} + x_9) + 5 \times (x_8 + x_7 + x_6) + 7 \times (x_5 + x_4 + x_3) + 9 \times (x_2 + x_1)$
となる，各ゴミの位置の累積和を前処理で $O(n)$ で計算すると， $O(n / k)$ 回の加算で求めることができる．全体の計算時間は調和級数で $n \times \sum_{k = 1}^{n} 1 / k = O(n \log n)$ となる．

この問題はlong long型で計算しても途中でオーバーフローする恐れがあるので，上の分割のコストの計算で $k$ が大きい順に計算して，分割の計算の途中で現在の最良値よりも悪くなるなら打ち切るということをしている．

計算時間： $O(n \log n)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    ll n, X;

    cin >> n >> X;
    vector<ll> x(n + 1);
    for (int i = 1; i <= n; ++i) cin >> x[i];
    for (int i = 1; i <= n; ++i) x[i] += x[i - 1];

    ll min_e = numeric_limits<ll>::max();
    for (ll k = n; 1 <= k; --k) {
        ll sum = 0, id = 1;
        for (ll r = n, l; 1 <= r && (sum + k * X) < min_e; r = l - 1) {
            l = max(1ll, r - k + 1);
            sum += (id == 1 ? 5 : 2 * id + 1) * (x[r] - x[l - 1]);
            ++id;
        }
        min_e = min(min_e, sum + k * X);
    }

    cout << min_e + n * X << endl;

    return 0;
}