ABC441Eを解いた【部分列の数え上げ】

筆者はレート800前後の茶～緑コーダ

ABC441のE問題を解いていく

実装コード

文字列を左から走査し、各位置で終わる A と B からなる部分列を累積して数え上げる。

• cnt は、これまでに各バランスの値が現れた回数を管理する配列にする
• i は現在のバランスを表し、負の添字を避けるために最初は N に置く
• t は、現在位置までで数えた条件を満たす部分列の数を表す
• 文字が A のときは、現在のバランス i が現れた回数を t に加え、i を 1 増やす
• 文字が B のときは、先に i を 1 減らし、そのバランスが現れた回数を t から引く
• 現在のバランス i の出現回数を cnt[i] に加える
• 各文字を処理したあとの t を ans に加算し、すべての位置で終わる部分列の数を合計する

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

def main():
    N = rI()
    S = rS()
    
    cnt = [0] * (2*N+1)
    t = 0
    ans = 0
    i = N
    cnt[i]+=1
    for c in S:
        if c == "A":
            t += cnt[i]
            i += 1
        elif c == "B":
            i -= 1
            t -= cnt[i]
        cnt[i] += 1
        ans += t
    print(ans)
if __name__ == '__main__':
    main()

感想

部分列の数え上げは、走査しながら途中までの個数を更新していく考え方が累積和に似ていると感じた。
文字を A と B に対応する増減として見ると、状態を整理しやすかった。