ABC418-D を Python で解いた

ABC418 D - XNOR Operation
https://atcoder.jp/contests/abc418/tasks/abc418_d

公式解説では DP で解いていますが僕は思いつけず、代わりに差分を更新していく方針で解けました。考え方としては尺取り法に近いかもしれません。

コンテスト中の動き

350点のC問題をじっくり考えて20分ほどかけて正解し、D問題もじっくり考えて23分ほどで正解できました。考察がうまくいきました。その後E問題も同様に考察がうまくいき、終了15分前に5完できました。

終わってみればパフォーマンス・順位ともに過去最高、レートも70増えて3桁から1000代半ばに戻すことができました。最高の結果です。

考察

本番中に考えていたことを書いていきます。

0の数の法則に気づく

まずは第一段階として0の数の法則に気づく必要があります。入力例1を見ながら考えてみると

とりあえず最後に "10" になってしまうと駄目そうなのはわかりますね。そこから "1" にすることはできません。ここで「最後の状態から逆算する」という方針を思いつきます。操作を1回するたびに文字が1つずつ減っていきますから、"1" になる直前は "00" または "11" です。次にそうなる直前は……と考えようとして、これだと無限に状態が拡がっていくだけで法則らしきものは見えてきません。この方針は捨てます。

もう少し考えてみると連続する 1 は 最終的に 1 つの 1 にできることに気づきます。また、0 を消す方法は "00" から "1" にする以外に無いことに気づきます。つまり連続する 2 つの 0 を見つけた場合にそれらを消していくことができます。

次に "01", "10" が "0" にできることから、"0111111" とか "11110" みたいな形があっても間の 1 を全部消して "00" にできることに気づきます。ということは "0111110" → "00" → "1" のように 2 つの 0 を見つけて、その間にある 1 を全部消して 0 同士が隣接するようにし、"00" を作ることで 0 を消すことができます。文字列から（間に挟まれる 1 を無視して）隣り合う 2 つの 0 を見つけ出し、2 つペアにして 1 組ずつ消していくイメージですね。となると文字列に含まれる 0 が偶数個なら美しい文字列となり、奇数個なら最後にぼっちになる 0 が残って "1" を作れなくなることがわかります。いわゆる popcount の考え方になりそうですね。

美しい文字列の数え方を考える

第二段階として、じゃあそれをどうやって数えるのか？という問題になります。全てのパターンを探索していては当然 TLE になるので工夫が必要です。この問題では 0 の個数の偶奇を調べるわけですが、偶奇を考えるときは累積和を使えば経験上うまく行くことが多い気がします。累積和といっても、個数ではなくその偶奇を記録していきます。0の個数を調べるのではなく偶と奇の2通りだけに注目することで何かと考えやすくなるというやつです。用意された入力例だとちょっと考えづらいので自分で例を作って考えてみましょう。T = "01101" の場合を考えてみます。

部分文字列を考える必要があるので 2 文字目から始まる場合も考えましょう。

同様に左端を動かしながら全ての場合を書き出してみます。

当たり前ですが、0 が来るたびに偶と奇が入れ替わりますね。そして一番左の状態はそれぞれ T[1] から T[5] の値によって決まります。また、T[2] から見る場合と T[3] から見る場合、それから T[3] から見る場合と T[4] から見る場合を比べると 偶偶奇奇 → 偶奇奇 → 奇奇 と左端が 1 つずつ消えていくだけで後ろは同じになることに気づきます。この表を眺めているうちに、最初に T[1] から見る場合を全て調べておいて、あとは差分更新で済ませられないか？ということを思いつきます。

まず T[1] から見る場合は 奇奇奇偶偶 となり、奇数になる場合が 3 個、偶数になる場合が 2 個です。ここから T[1] を削除することで T[2] から見る場合を作ります。左端の 1 つを削除するほか、T[1] = 0 ですから以後の偶奇が入れ替わります。なので奇数になる場合が 2 個、偶数になる場合が 2 個です。ここから T[2] を削除することで T[3] から見る場合を作ります。T[2] = 1 なので左端の 1 つを削除するだけで十分です。要するに偶数になる場合が 1 減ります。以後同様にやっていきます。

このやり方でいけそうですね。あとはこの操作を実装するだけです。差分更新するたびに偶の数を足していき、その合計数が問題の答えとなります。

実装

N = int(input())
T = "." + input() # 1-index にして扱いやすくする
# 最初は左端から右端まで全部調べる。個数の偶奇の累積和にする。偶数なら 0, 奇数なら 1 とする。
cum = [0 for _ in range(N+1)]
for i in range(1, N+1):
  if T[i] == "0":
    cum[i] = 1 - cum[i-1]
  else:
    cum[i] = cum[i-1]
# 次に差分更新をする。
even, odd = N - sum(cum), sum(cum) # 初期値
count = even # 初期値
for i in range(2, N+1):
  if T[i-1] == "1":
    even -= 1
  elif T[i-1] == "0":
    odd -= 1
    even, odd = odd, even
  count += even
print(count)