ABC270E - Apple Baskets on Circle
Published on 2023-09-26Last Modified 2023-09-28
Table Of Contents
問題概要
円環状に並べてある$1$から$N$の番号が付いた$N$個のかごがある。 かごは$1$から$N$まで順番に並んでおり、かご$N$の隣にはかご$1$がある。
かご$i$にはりんごが$A_i$個入っている。
高橋君は、以下の行動を繰り返す。
- 目の前にあるかごの中にりんごがあれば$1$個かごから取り出して食べる。その後、りんごを食べたかどうかにかかわらず、右隣のかごの前に移動する。
高橋君がちょうど$K$個のりんごを食べた時点で、各かごに入っているりんごの数を求めよ。
制約
- $1 \leq{} N \leq{} 10^5$
- $0 \leq{} A_i \leq{} 10^{12}$
- $1 \leq{} K \leq{} 10^{12}$
- $K \leq{} \sum A_i$
解法
$K$の値が非常に大きいため、高橋君の動きをシミュレートすることはできない。 そこで、高橋君の動きを少しずつまとめよう。
例えば、ある時点において$0$個を除いたりんごの最小値が$x$個であって、かつ$1$個以上のりんごが入ったかごが$y$個であったとする。
この時、少なくとも$x$周している間は$y$は一定になるはずであるから、素直にシミュレートすると$x$周必要なところを$1$周にまとめることができる。 $x$周しても$A_i$における大小関係は逆転しない($A_i < A_j$が$A_i = A_j$になることはあるが、$A_i < A_j$が$A_j < A_i$となることはない)ので、次の最小値$x$を容易に計算することができる。
また、$x$周している間に$K$個を超えてしまうときは、$K$個を超えないで何周できるかを計算するとよい。 具体的には、今まで食べたりんごを$\mathrm{sum}$として、$\lfloor{} (K-\mathrm{sum})/y \rfloor{}$とすればよい。
最後の一周は高々$N$回の操作で終わるので、シミュレートすればよいだろう。
そこで、次の解法を得る。
解法1
詳しいアルゴリズムは言語による説明よりもソースを見る方が早いかと思うので、D言語による実装を掲載する。
Remainingはその時点で$0$個以上のりんごが入っているかごの数である。
また、現時点で最小のりんごの数を得るために優先度付きキューを利用している。
これは上で説明した大小関係が保存されることを利用している。
void solve (int N, long K, long[] A) {
long sum = 0;
long Remaining = 0;
foreach (a; A) if (0 < a) Remaining++;
BinaryHeap!(pair[], "b.val < a.val") PQ = [];
foreach (i, a; A) PQ.insert(pair(cast(int) i, a));
while (sum < K) {
auto head = PQ.front; PQ.removeFront;
with (head) {
if (A[idx] == 0) continue;
if (K <= sum + Remaining) {
// 一周とれば終了
int i = 0;
while (sum < K) {
if (0 < A[i]) sum++, A[i]--;
i++;
}
break;
}
if (sum + Remaining < K) {
// A[idx]が最小なので、これを上限にして、できるだけとる
long take = min(A[idx], (K-sum)/Remaining);
sum += Remaining*take;
foreach (ref a; A) if (0 < a) {
a -= take;
if (a == 0) Remaining--;
}
}
}
}
// output
foreach (i; 0..A.length) write(A[i], (i == A.length-1 ? "\n" : " "));
}
この解法は確かに正しい解を得るが、実は$\text{worst} ~ O(N^2)$となっている。(筆者はこれに気づかずに2TLEした。) 例えば$A_i = i$であって、$K = \sum A_i$であるときがこのケースに当たる。 更新が$N$行われて、かつ更新一回で$O(N)$回の操作が必要であるからだ。
素直にシミュレートするよりはかなり高速化したが、これでは間に合わない。
解法2
さて、少し突飛な発想であるかもしれないが、$x$周したときに何個りんごをとれるかを考えよう。 $x$個に満たないかごからは$A_i$個までしかとれないので、$\sum \min{} (x, A_i)$個ということになる。
この値は$x$に対して(適切な区間で)狭義単調増加する。
すると、この値が$K$を超えるかどうかで二分探索ができることが分かる。 要するに、何周までなら$K$を超えないのかを高速に求めることができる。
最後の一周は素直にシミュレートすればよいので、これで解ける。 時間計算量は$O(N \log{} (\max{} A_i))$である。
以下にD言語による実装例を載せる。
void solve2 (int N, long K, long[] A) {
// A[i]からとれるならX個とるとするとき、その総和がKを超えるかどうかで二分探索
// 二分探索: f(x) := sum( min(x, A[i]) ) に対して、f(ok) <= K < f(ng)
// ok <- [0, max(A)]
long f (long x) {
long sum = 0;
foreach (a; A) sum += min(x, a);
return sum <= K;
}
long ok = 0, ng = 10L^^12+1;
while (1 < abs(ok-ng)) {
long mid = (ok+ng)/2;
if (f(mid)) {
ok = mid;
} else {
ng = mid;
}
}
long sum = 0;
// 少なくとも(とれるなら)ok個とってよい。
foreach (ref a; A) sum += min(a, ok), a -= min(a, ok);
// 端数を一周して合わせる
{
int i = 0;
while (sum < K) {
if (0 < A[i]) sum++, A[i]--;
i++;
}
}
// output
foreach (i; 0..A.length) write(A[i], (i == A.length-1 ? "\n" : " "));
}
振り返り
解法2は解法1をさらにまとめた形だということもできるだろう。 筆者がこの解法に至るまでに次のような手順を踏んだ。
- (解法1がTLEして)$A$がばらばらの値の時にまずいのか…
- なんか何回も同じ要素を引き算してんな…どうにかならんかな…
- $0$個になったかごは無視してもいいから…(天啓が下りてきて)可能なら$x$個とる方針で行けば二分探索できそう？
…うまく言語化できない。というよりは、どうやって思いついたのか詳細にはわからないというのが正しいのか…
ただ、やはり最後の一周とそれまでを区別して考えるのは大事なポイントだと思う。
もし各かごにりんごが無限に入っていたとしたらこのように「何周するのか」に着目する解法を容易に構成できると思う。 が、今回は、りんごの減少に伴って変化が生じるので発想の難易度が上がっていると思った。
より易しい問題設定を考えたり、過去に解いた問題の記憶をためることが重要なのかなと思う。
(2023-09-28追記)
深夜に考えていたら少し思いついたので補足します。 この手の問題は「解の構造」を考えることが大事だと思います。
K個目のりんごを食べるときにかごのりんごがどのように減っているかを考えると、 0になるまで減る、または今まで周回した分だけ減るということが分かります。
すると、何周したかに着目できるのかなと思いました。 何周したかに着目できれば二分探索に帰着するのはそんなに難しくないと思います。
しかし、この手の問題は大抵シミュレーションの高速化で解けることが多いので、解法1を組み切る力も必要かなと思います。 (もちろん、シミュレーションの最悪計算量を一発で見抜く力も)