ARC137B - Count 1's

考え方

部分列の中含まれる「(1の数) - (0の数)」のパターンが何通りあるか，という問題．

累積和のパターンが何通りあるかを考えれば良さそうだが，1の数が変化しても，0の数が変化しても累積和が変化してほしい．そこで，0を-1で置き換えて累積和を考える．

区間$L_{1} = [a_{1}, a_{2}]$で累積和が最小になり，区間$L_{2} = [b_{1}, b_{2}]$で累積和が最大になるとする．このとき，$L_{1}$から$L_{2}$へは，両端を1ずつ増減させることで変形できる．両端を1ずつ増減させたときの累積和の変化量は$\pm 1$なので，区間を$L_{1}$から$L_{2}$へ変形させる仮定で，その間の値を全て取ることになる．

よって，累積和の最大値・最小値を求めれば答えは「(最大値) - (最小値) + 1」となる．

全ての区間の累積和を計算するにはは，右端と左端を決める必要がある．これを全通り試すのは間に合わない．しかし，知りたいのは最大値と最小値だけなので，右端が$R$のときの累積和の最大値・最小値は，それぞれ

\begin{aligned}
\sum_{i\in [0,R]} a_{i} - \min \Biggl\{\sum_{i\in [0,L]} a_{i} \biggl| L = 0,1,...,R\Biggr\}
\end{aligned}

および

\begin{aligned}
\sum_{i\in [0,R]} a_{i} - \max \Biggl\{\sum_{i\in [0,L]} a_{i} \biggl| L = 0,1,...,R\Biggr\}
\end{aligned}

と計算できる．つまり，右端を一つずつずらしながら（左端がゼロの）累積和を計算していき，その時点の（左端がゼロの）累積和の最大値・最小値を保持しておけばよい．

回答例

N = int(input())
A = list(map(int, input().split()))
A = [1 if a == 1 else -1 for a in A]

acc = 0
acc_0R_max, acc_0R_min = 0, 0
acc_LR_max, acc_LR_min = 0, 0
for a in A:
    acc += a
    acc_0R_max = max(acc_0R_max, acc)
    acc_0R_min = min(acc_0R_min, acc)
    acc_LR_max = max(acc_LR_max, acc - acc_0R_min) 
    acc_LR_min = min(acc_LR_min, acc - acc_0R_max)

print(acc_LR_max - acc_LR_min + 1)