この時代、この手の Greedy はたくさんあったのね!
問題概要
長さ の数列 が与えられる。「数列の要素を 1 つ選び、+1 するか、-1 する」という操作を行うことで、次の状態を達成したい。ただし、 とする。
- すべての に対して、 である
- すべての に対して、 と の符号は異なる
制約
考えたこと
まず、先頭要素を正にする場合と、負にする場合とで、その後の状況が大きく異なるので、場合分けして考えることにしよう。
先頭要素を正にする場合さえ解ければ、負にする場合も同様に解けると思われるので、先頭要素を正にする場合のみ考えることにした。
先頭要素が 0 以下であるとき
この場合は、先頭要素を 1 にすれば良い。1 より大きくする必要はないことに注意しよう。なぜならば、1 より大きくすると
- 先頭要素を 1 より大きくするコスト自体が、1 にするコストよりも大きい
- さらに、その後 2 番目の要素を変更することで「先頭要素 + 2 番目の要素」を負にする必要があるのだが、そのコストも増大してしまう(より正確には、「先頭要素 + 2 番目の要素」を -1, -2, ... のどの値にするとしても、そのためのコストが大きくなる)
ということで、デメリットしかないからだ。よって、先頭要素が負値である場合は、先頭要素を 1 にすればよい。
先頭要素が正の値であるとき
この場合は、先頭要素を 1 まで下げておくかどうかが悩ましい。しかし、実は先頭要素に何もしなくてよいことが言える。
たとえば、先頭要素が 5、2 番目の要素が 3 だったとする。このとき、
- 先頭要素を 5 のまま、先頭と 2 番目の要素の和を -1, -2, ... にするためのコストは、-9, -10, ... である
- 先頭要素を 1 にした上で、先頭と 2 番目の要素の和を -1, -2, ... にするためのコストも、-9, -10, ... である
というわけで、実は「先頭要素で値を下げず、その分を 2 番目の要素の値を下げる」ことにしても、コストは変わらないのだ。
逆に、2 番目の要素がもともと負値である場合、先頭要素を 1 に下げることにデメリットが生じることもある。たとえば、先頭要素が 5、2 番目の要素が -6 である場合、何もしなくても条件を満たしている。しかし、なまじ先頭要素を 1 にしてしまうと、その分コストが生じてしまうのだ。
まとめ
以上の考察をまとめると、次の解法にたどり着く。
先頭要素を正にする場合と負にする場合については両方試す。
一般に、 の符号を決めた状態で、 への操作を考えるときを考える。
- を正の値にしたいとき:
- の場合は、その値が 1 になるまで の値を増やせば良い
- の場合は、何もしなくてよい
- を負の値にしたいとき:
- の場合は、その値が -1 になるまで の値を減らせば良い
- の場合は、何もしなくてよい
以上の解法の計算量は となる。
コード
#include <bits/stdc++.h> using namespace std; int main() { int N; cin >> N; vector<long long> A(N); for (int i = 0; i < N; i++) cin >> A[i]; long long plus_cost = 0, sign = 1, sum = 0; for (int i = 0; i < N; i++) { sum += A[i]; if (sign == 1) { if (sum <= 0) plus_cost += 1 - sum, sum = 1; } else { if (sum >= 0) plus_cost += sum + 1, sum = -1; } sign *= -1; } long long minus_cost = 0; sign = -1, sum = 0; for (int i = 0; i < N; i++) { sum += A[i]; if (sign == 1) { if (sum <= 0) minus_cost += 1 - sum, sum = 1; } else { if (sum >= 0) minus_cost += sum + 1, sum = -1; } sign *= -1; } cout << min(plus_cost, minus_cost) << endl; }