Atcoder解いてみた ABC 059 C - Sequence

ABC 059 C - Sequence 解いてみた

問題はこちら
こちらの方が、スマートに解いていました。

方針

配列の最初の要素から値を更新していき、1つずつ条件*を満たすようにしていく。
→ ランダムな順番で要素の値を更新しても、1番目から順番に更新していっても、更新回数は変わらない。

条件*

a : 第 1 項から第 i 項までの和は 0 でない。
b : i 項までの和と i+1 項までの和の符号が異なる。

結局、以下のように、実装すればよい。

1. 最初(0番目)の要素が正と仮定して、偶数番目までの和が1以上、奇数番目までの和が-1以下になるように配列値を更新していく。更新回数は記録
2. 最初(0番目)の要素が負と仮定して、偶数番目までの和が-1以下、奇数番目までの和が1以上になるように配列値を更新していく。更新回数は記録
3. 1,2で求めた更新回数の小さい方を答えとして出す。

注意

問題文に記載の通り、+-1の値で少しずつ値を更新して、更新回数を一回ずつ記録しても、時間オーバーになってしまうため、以下ような工夫点が必要
例) 値を-2から1に変更したいとき、更新回数をabs(-2)+1で求める。

実装の都合上、配列内の値を直接更新するのではなく、i番目までの合計値に対して、何回、更新を行えばいいのか求めるように実装した。

回答

#include <iostream>
#include <cmath>
#include <algorithm>
 
using namespace std;
 
long long solve(long long* arr, int n, bool first_plus);
 
 
 
int main() {
	long long a[100000];
	int n;
 
	cin>>n;
	for (int i=0; i<n; i++) {
		cin >> a[i];
	}
 
	long long min = 100000;
	long long ans_1 = 0;
	long long ans_2 = 0;
 
	ans_1 = solve(a, n, true);
	ans_2 = solve(a, n, false);
 
	std::cout << std::min(ans_1, ans_2) << std::endl;
 
	return 0;
}
 
// first_plus : true →最初の数がプラス
// first_plus : false →最初の数がマイナス

long long solve(long long* arr, int n, bool first_plus) {
 
	long long count = 0;
	long long tmp_sum = 0;
 
	for (int i = 0; i<n; i++) {
		tmp_sum += arr[i];
 
			if (first_plus == true && i % 2 == 0 && tmp_sum <= 0) {
				count += abs(tmp_sum) + 1;
				tmp_sum = 1;   //count回の操作を行った後、合計値は、1になることが確定
			}
			else if (first_plus == true && i % 2 == 1 && tmp_sum >= 0) {
				count += tmp_sum + 1;
				tmp_sum = -1;  //count回の操作を行った後、合計値は、-1になることが確定
			}
			else if (first_plus == false && i % 2 == 0 && tmp_sum >= 0) {
				count += tmp_sum + 1;
				tmp_sum = -1;  //count回の操作を行った後、合計値は、-1になることが確定
			}
			else if (first_plus == false && i % 2 == 1 && tmp_sum <= 0) {
				count += abs(tmp_sum) + 1;
				tmp_sum = 1;   //count回の操作を行った後、合計値は、1になることが確定
			}
	}

	return count;
}