あるエンジニアのAtCoder奮闘記

東京都港区にあるアミフィアブル株式会社のエンジニアが、AtCoderで解いた問題について振り返ったりしていく会社公認のブログです。

AtCoder AGC 034 B - ABC (Python)

irisruneです。当初Pythonで解けなくてGoで解いていたのですが、ミスに気付いたのでPythonで解き直しました。TLEしたときにすぐPythonのせいにするのはよくないですね。

atcoder.jp

600点としてはかなり簡単な方で、少なくとも土曜日の企業コン500点問題よりは簡単だと思いました。

s = input()
ans = 0
cntA = 0
for i in range(0, len(s) - 1):
    if s[i] == 'A':
        cntA += 1
    elif s[i: i + 2] == "BC":
        ans += cntA
    elif i > 0 and s[i - 1: i + 1] == "BC":
        continue
    else:
        cntA = 0
print(ans)

方針としては、

  • "A"と"BC"が隣接していると"ABC"→"BCA"となる。
  • "AAA...A"と連続しているケースについても上記の操作を繰り返すと、"AAA...ABC"→"BCAAA...A"となる。
  • 同様に"BCBCBC...BC"と連続しているケースについても上記の操作を繰り返すと、"ABCBCBC...BC"→"BCBCBC...BCA"となる。

という規則性を見つければ"A"と"BC"とそれ以外("AC"とか"BB"とか)を判別しながら処理を行うだけとなります。"BC"が出現する直前の連続する"A"の個数が変換後の末尾の"A"の個数と一致するのがポイントですね。

ただ後で気付いたので仕方ないですが、解説のように"BC"を"D"に置換した方がかなり解きやすいですね。置換しなかった場合は上記コードのようにインデックスで回しながら隣接文字と同時にチェックするか、フラグ管理しながら1文字ずつチェックするかになってどちらにせよ煩雑になります。

雑記

  • Pythonで最初解こうとしたときは規則性を把握しきれず、二重ループにしてしまい計算量がO(N^2)になってTLEしてしまいました。
    • 慣れていない言語で書くと色々気が回らなくなるものですね。
  • Goで"BC"を"D"に置換するのはそれはそれで難しいような気もしなくはないです。
    • 文字列分割→文字列結合という処理になりそうですが、文字列長が大きいときに計算量が多くかかる気がします。
    • stringsパッケージを使えば置換ができるようなのであとは計算量次第でしょうか。
      • 追記:置換を行ってもほとんど実行時間変わらずにACできました。