けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

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

AtCoder AtCoder300点 ABC-C 水色diff NoviSteps2Q Greedy Greedy:ある量を決めると残りが決まっていく Greedy:端から順に決まっていく 数列 負値が本質的に効く問題 操作 操作:±1 禁止文字列 最小回数・最小個数を求める 最適化問題 場合分け ARC-C

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

問題概要

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

  • すべての i に対して、 B_{i} \neq 0 である
  • すべての i に対して、 B_{i} と B_{i+1} の符号は異なる

制約

  • 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} \le 0 の場合は、その値が 1 になるまで A_{i} の値を増やせば良い
    • A_{1} + \dots + A_{i-1} + A_{i} \ge 1 の場合は、何もしなくてよい
  • A_{1} + \dots + A_{i-1} + A_{i} を負の値にしたいとき:
    • A_{1} + \dots + A_{i-1} + A_{i} \ge 0 の場合は、その値が -1 になるまで A_{i} の値を減らせば良い
    • A_{1} + \dots + A_{i-1} + A_{i} \le -1 の場合は、何もしなくてよい

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

コード

#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;
}