programming

AGC040 A - >< の備忘録

わりと頻出しそうなアルゴリズムなので記録

You can share this article with

問題

問題は公式を参照してください。

最初に思いついた内容

左から順に<が来たら+1して>が来たら-1した配列lsを作る。
が、そのままだと非負数列なのに負の数が出現してしまう。そこで、min(ls)を各項から引いてから和を取る。

しかし、それでもダメで、例えば<>>のときに前処理で[0,1,0,-1]が出来る。
ここからmin(ls)である-1を各項から引くと、[1,2,1,0]となってしまう。不等号の制約は満たすが、最小でない。ls[0]は0でなくてはならない。

次に思いついた内容

左右両側から見て成り立っていることが重要っぽい。
1 度の操作で両方やろうとすると頭がバグるので、片側ずつの単調増加の操作に分解する。

不等号が連続する場合にコンボがつながっていくが、 n個の小なりが連続した場合は、x {<*n} x+nのようになる。 逆にn個の大なりが連続した場合は、x=n {>*n} xのようになる。

各インデックスに対し、左にある小なりの連続数と右にある小なりの連続数を前処理で記録しておき、max を取ればよい。 この時点で公式の解説 PDF とほぼ同じ考え。

S = input()
# 左側から<の連続数をカウント
count = 0
ge = [0]
for i, v in enumerate(S):
    if v == '<':
        count += 1
    else:
        count = 0
    ge += (count,)
count = 0
# 右側から>の連続数をカウント
le = [0]
for i, v in enumerate(S[::-1]):
    if v == '>':
        count += 1
    else:
        count = 0
    # appendleftは計算量O(N)なのでダメ
    le += (count,)
ls = []
# リストleは逆順のため反転する
for g, l in zip(ge, le[::-1]):
    ls += (max(g, l),)
print(sum(ls))

ちなみにappendleftがO(1)のdequeも、任意箇所からの取り出しは最悪O(N)で遅いのでダメ。

最終的にたどり着いた答え

最後に追加で for ループ回すのがちょっとセンスない感じがしたので、ここを削りたい。 左から見たときは単純に<なら+1する操作をしていけばよい。 右から見るときだけ最後の for ループでやっていたmaxを取れば同じことができる。

A0<A1<A2>⋯<AnA_0<A_1<A_2>\dots<A_n

みたいなイメージで、各不等号がSiS_iだとすれば、

A0 S0 A1 S1 A2 S2…Sn−1 AnA_0\ S_0\ A_1\ S_1\ A_2\ S_2\dots S_{n-1}\ A_n

という並びでインデックスを取っていく。
最終的には以下のコードになった。

s = input()
a = [0] * (len(s) + 1)
# a[i]とa[i+1]の間がs[i]とすれば
for i in range(len(s)):
    if s[i] == '<':
        a[i + 1] = a[i] + 1
for i in range(len(s)):
    if s[-(i + 1)] == '>':
        a[-(i + 2)] = max(a[-(i + 1)] + 1, a[-(i + 2)])
print(sum(a))

ちょっとインデックスをバグらせやすい気がするのでもう少し洗練させたいが、とりあえず計算量落とすならこれかな、というところでこの記事はおしまい。