与えられた条件を満たすために必要な最小操作回数を求めるC++プログラム

N個の要素からなる配列Aがあるとします。1回の操作ごとに、任意の要素を1つ選び、その値を1増やすか1減らすことができます。ここで、次の2つの条件を満たすために必要な最小の操作回数を求めるのが目標です。

• 1以上n以下のすべてのiについて、第1項から第i項までの総和(接頭辞和)が0にならないこと。
• 1以上n−1以下のすべてのiについて、第1項から第i項までの総和の符号が、第1項から第(i+1)項までの総和の符号と異なること。

言い換えると、接頭辞和は正と負の間で交互に入れ替わりながら推移し、途中で一度も0になってはいけない、という条件です。

入力例と出力例

たとえば入力が A = [1, -3, 1, 0] のとき、答えは4になります。4回の操作で配列を 1, -2, 2, -2 に変形でき、このとき第1項〜第4項までの累積和はそれぞれ 1, -1, 1, -1 となり、両方の条件を満たすからです。

解き方のステップ

この問題は貪欲法で効率よく解けます。接頭辞和の符号の並び方は「正・負・正・負…」と「負・正・負・正…」の2通りしかありません。そこで、両方のパターンについて配列を左から順に走査し、条件を満たすようにその場で要素を補正したときの合計コストを求め、小さい方を答えとします。

1. 現在の接頭辞和sumに次の要素aiを加えた値をnsumとします。
2. 期待する符号が正なのにnsum ≤ 0のときは、aiに|nsum| + 1を加算して接頭辞和を1以上にします。このときの操作回数は|nsum| + 1回です。
3. 期待する符号が負なのにnsum ≥ 0のときは、aiからnsum + 1を減算して接頭辞和を−1以下にします。操作回数はnsum + 1回です。
4. 各要素について処理しながらsumと符号フラグsを更新し、最後に合計retを返します。

n := size of A
ret := 0
sum := 0
for each ai in A, do
   nsum := sum + ai
   if s > 0, then:
      if nsum <= 0, then:
         ret := ret + |nsum| + 1
         ai := ai + |nsum| + 1
   otherwise
      if nsum >= 0, then:
         ret := ret + nsum + 1
         ai := ai - (nsum + 1)
   sum := sum + ai
   s := s * (-1)
return ret

C++による実装例

理解を深めるために、以下の実装を見てみましょう。

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

int util(vector<int> A, int s){
   int n = A.size();
   int ret = 0;
   int sum = 0;
   for (int ai : A){
      int nsum = sum + ai;
      if (s > 0){
         if (nsum <= 0){
            ret += abs(nsum) + 1;
            ai = ai + abs(nsum) + 1;
         }
      } else{
         if (nsum >= 0){
            ret += nsum + 1;
            ai = ai - (nsum + 1);
         }
      }
      sum += ai;
      s *= -1;
   }
   return ret;
}
int solve(vector<int> A){
   int res = min(util(A, 1), util(A, -1));
   return res;
}
int main(){
   vector<int> A = { 1, -3, 1, 0 };
   cout << solve(A) << endl;
}

入力

{ 1, -3, 1, 0 }

出力

4

計算量

util関数は配列を一度だけ走査するため、時間計算量はO(N)です。solve関数で2つの符号パターンを試しても、全体の計算量はO(N)のまま変わりません。なお、ベクトルを値渡ししているため、追加で必要となる空間計算量はO(N)となります。