AGC027 B. Garbage Collector
問題
数直線上にN個のゴミが落ちており、左からi番目のゴミの位置はx_iにある。
ゴミ箱は位置0にあり、掃除ロボットを動かすことですべてのゴミをゴミ箱に入れたい。
掃除ロボットには以下の制約がある。
- 最初は位置0にいる
- 現在持っているゴミの数をkとすると、距離1あたりの移動には(k+1)^2だけエネルギーがかかる
- ゴミを拾う、または、ゴミを捨てるのにはXのエネルギーを消費する
すべてのゴミを回収するのにかかるエネルギーの最小値を答えよ。
制約
- 1 <= N <= 2 * 10^5
- 0 < x_i <= 10^9
- x_iは整数
- x_i < x_{i+1}
- 1 <= X <= 10^9
解法
ロボットの動きを考えてみると、行きは何も回収せず、帰りがけにゴミをいくつか回収して戻ってくるのがエネルギーを抑えられそうというのがわかる。
ただ、どのゴミを回収すればよいかはぱっと見ではわからない。
また、移動エネルギーは持っているゴミの数に依存するので、「あるゴミの回収にかかるエネルギー」は独立に計算できない(そのゴミ以外に同時に回収するゴミがなんなのかが計算には必要)。
そこで、まず、あるk個のゴミを考え、それを回収するのにかかるエネルギーがどのように計算されるか考える。
拾う・捨てるのエネルギーは、各ゴミを拾う「k * X」と捨てる「X」の和「k * X + X」になっている。
今、ゴミの位置が左から「a,b,c,d」だとすると、移動にかかるエネルギーを考えると、
行き：1 * d
帰り：4 * (d-c) + 9 * (c-b) + 16 * (b-a) + 25 * a
なので、これを足して整理すると、
5 * d + 5 * c + 7 * b + 9 * a
となり、「ゴミ集合の遠い方のゴミから位置に5,5,7,9,11,13,...をかけたもの」になっていることがわかる。
したがって、N個のゴミをいくつかの集合に分解して考えるとき、その集合の移動エネルギーは上記の方法で求められる。
N個のゴミをm個の集合に分けることを考えると、上の「遠い方のゴミから位置に5,5,7,9,11,13,...をかけたもの」がその集合の移動にかかるエネルギーになることから、できるだけ遠くの位置にあるゴミに小さな係数がかかるように集合を定義してあげたほうがエネルギーが小さくなり良いということがわかる。
これは、例えば、m=3で、ゴミの位置が左から「a,b,c,d,e,f,g,h,i,j,k,l」のような場合だと、(m=)3つの各集合は、(遠い方から)
- l, i, f, c
- k, h, e, b
- j,g,d,a
のように選んであげれば、
- 5 * l + 5 * i + 7 * f + 9 * c
- 5 * k + 5 * h + 7 * e + 9 * b
- 5 * j + 5 * g + 7 * d + 9 * a
がそれぞれの集合での移動にかかるエネルギーになっている。
これは、位置が遠い方からm個ごとのかたまりを作って、遠い塊から5,5,7,9,11,13,...をかけている。
m個の塊の合計は累積和を使えばO(1)で求められ、5,5,7,9,11,13,...をかける行為はceil(N/m)回の掛け算になっている。
また、拾う・捨てるのエネルギーは、各ゴミを拾う「N * X」と捨てる「m * X」の和「N * X + m * X」になっている。
mを1〜Nまでループをまわすことを考えると、各集合のエネルギーを求める部分でceil(N/m)回の掛け算を行うことになるが、「O(Σ_{1 <= i <= N} N/i) = O(NlogN)」である性質を考えると、2重ループを回しても全体でO(NlogN)の計算量になる。
桁あふれがあり得るので溢れないように注意。
