AGC040 A - >< の備忘録
わりと頻出しそうなアルゴリズムなので記録
問題
問題は公式を参照してください。
最初に思いついた内容
左から順に<
が来たら+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
を取れば同じことができる。
みたいなイメージで、各不等号がだとすれば、
という並びでインデックスを取っていく。
最終的には以下のコードになった。
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))
ちょっとインデックスをバグらせやすい気がするのでもう少し洗練させたいが、とりあえず計算量落とすならこれかな、というところでこの記事はおしまい。