AtCoder ABC 059 C - Sequence (ARC 072 C) (2Q, 水色, 300 点) (original) (raw)

この時代、この手の Greedy はたくさんあったのね！

問題へのリンク

問題概要

長さ $N$ の数列 $A_{1}, A_{2}, \dots, A_{N}$ が与えられる。「数列の要素を 1 つ選び、+1 するか、-1 する」という操作を行うことで、次の状態を達成したい。ただし、 $B_{i} = A_{1} + \dots + A_{i}$ とする。

制約

$2 \le N \le 10^{5}$

考えたこと

まず、先頭要素を正にする場合と、負にする場合とで、その後の状況が大きく異なるので、場合分けして考えることにしよう。

先頭要素を正にする場合さえ解ければ、負にする場合も同様に解けると思われるので、先頭要素を正にする場合のみ考えることにした。

先頭要素が 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 にしてしまうと、その分コストが生じてしまうのだ。

まとめ

以上の考察をまとめると、次の解法にたどり着く。

先頭要素を正にする場合と負にする場合については両方試す。

一般に、 $A_{1} + \dots + A_{i-1}$ の符号を決めた状態で、 $A_{i}$ への操作を考えるときを考える。

$A_{1} + \dots + A_{i-1} + A_{i}$ を正の値にしたいとき:
$A_{1} + \dots + A_{i-1} + A_{i}$ を負の値にしたいとき:

以上の解法の計算量は $O(N)$ となる。

コード

#include <bits/stdc++.h> using namespace std;

int main() { int N; cin >> N; vector 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;

}