【AtCoder】黒マスが連続する区間の数を数える

https://atcoder.jp/contests/abc411/tasks/abc411_c

問題

• クエリによって指示された1マスが黒く塗られる
• 再度塗られるクエリが来ると白に戻る
• 連続で塗られた黒マスの区間はいくつあるか、1つのクエリが処理されるごとに出力する

例

□□□□□　→　0
□□■□□　→　1
□□■□■　→　2
□□■■■　→　1
■■■■■　→　1

問題を解く上でのポイント

マスが塗られた後、全てのマスをチェックして区間の判定をしていると実行時間制限に引っかかるため、効率よく区間の判定を行う必要がある。

解答

ToDo：もうちょっとスマートに書きたい（自分が分かりやすいように書いたので、コードが汚い）

N, Q = map(int, input().split())
A = list(map(int, input().split()))

squares = [0] * (N + 2)
segments = 0
for i in range(Q):
    pos = A[i]
    
    # 色を反転
    squares[pos] = 1 - squares[pos]
    
    # 増えた
    if squares[pos] == 1 and squares[pos-1] == 0 and squares[pos+1] == 0:
        segments += 1
    if squares[pos] == 0 and squares[pos-1] == 1 and squares[pos+1] == 1:
        segments += 1
    # 減った
    elif squares[pos] == 0 and squares[pos-1] == 0 and squares[pos+1] == 0:
        segments -= 1
    elif squares[pos] == 1 and squares[pos-1] == 1 and squares[pos+1] == 1:
        segments -= 1
    
    print(segments)

以下順を追って解説

解説

解説1：マス目の初期化

黒区間判定の都合上、本来のマス目の左右両端に白マスが1マスずつ存在していると仮定する。

squares = [0] * (N + 2)

例えば5マス指定されたら両端に1マス追加して7マスとする。
こんな感じで
□|□□□□□|□←実際には存在しない7マス目

解説2:色の反転

squares[pos] = 1 - squares[pos]

この部分を実際に計算すると
□→■　1 - 0 = 1
■→□　1 - 1 = 0
となり色の反転ができている。
追記
排他的論理和（XOR）演算でもできる。（こっちの方がスマートやね…）

squares[pos] ^= 1

解説3：黒区間の変化パターン

クエリは1マスしか塗れず、変化は以下の3パターンしか存在しない。

1. 黒区間が1つ増える
2. 黒区間が1つ減る
3. 変わらない

黒区間の数が増減する1,2パターンでのみカウント数を増減させればよい。

解説4：黒区間の数が増減しているのはどんな時か？

具体的にどんな場合に黒区間が増減するのか考える。

解説4-1：増える場合

□□□□□　→　□□■□□
□■■■□　→　□■□■□

1. 塗ったマスが黒、かつ両隣が白（＝新しい黒区間ができる）
2. 塗ったマスが白、かつ両隣が黒（＝黒区間を分断して増やす）

解説4-1：減る場合

□■□■□　→　□□□■□
□■□■□　→　□■■■□

1. 塗ったマスが白、かつ両隣が白（＝1つの黒区間を消す）
2. 塗ったマスが黒、かつ両隣が黒（＝分断された黒区間を繋げる）

これ以外のパターンでは黒区間が増減することはない。
よってい以上の4パターンにおいて、区間カウンターを増減させればOK