Zenn
🐍

AtCoderのとある問題でCode Golfに挑戦した記録 (Python3)

2022/11/11に公開
Python
AtCoder
codegolf
tech

きっかけ

最近 @yanecoder さんが AtCoder の問題を解きまくっていて、さらにshortest記録更新にも熱心に取り組んでいるので、自分も以前に多少は嗜んでいたしと対抗心を燃やして挑んでみたところ、多くの学びがあったので記録として残しておく。

※使用言語はPython3です

Code Golf とは

とにかく短いコード(byte数)で目的を達成するプログラムを書くこと。

https://ja.wikipedia.org/wiki/コードゴルフ

AtCoder では AC をとれれば目的達成といえるし、言語別にすべての提出のコード長も表示されるので分かりやすい。

お題

本記事で取り上げる問題は、これ。

https://atcoder.jp/contests/iroha2019-day3/tasks/iroha2019_day3_e

5
/
\
/
/
\

のような入力が与えられて、「く」のように等しい長さの / と \ が連続しているものをカウントする、というもの。
上の例だと 2-3行目が / \ で等しいが、4-6行目は // \ と / の方が長いので対象外となり、正解は 1 となる。

単純に各文字のカウントしながら見ていくだけでは正解に辿り着けず、「/ がスタートした」「\ に折り返した」「また新しい / が始まった」などの状態遷移を管理しながら判定していく必要がある。

参考解答例 (134 Byte)

先月までの時点での、Python3の最短 AC コードはこれ。

https://atcoder.jp/contests/iroha2019-day3/submissions/10517249 (134 Byte by c_r_5)

_,*c=open(0)
a=b=s=0
for c in c:
  if c<'\\':
    if b:
      s+=a==b
      a,b=1,0
    else:a+=1
  else:b+=1
print(s+(a==b))

これでも十分に短くてすごいのだけど、ここからさらに短いコードを目指していく。

解読

まず入力の読み取りは input() を使わずに open(0) を使う。Golfでは常套テクニックのようだ。

https://qiita.com/neko_the_shadow/items/521d1361820c42547741

_,*c= で代入すれば、 c に2行目以降の文字列配列を受け取れるので、それを順番に見ていく。

a は / の連続する数をカウントし、 b は \ の連続する数をカウントするもの、のようだ。 s が対象の出現カウント。

c<'\\' すなわち / の行であれば b を判定し、 b が正かつ a と同値ならば s を増加させる。 s += (a==b) は比較結果のbool値を暗黙的に数値として += することで True のときだけ +1 するという意味になる。
b の真偽判定をする場面は \ が少なくとも1つ以上出現した後の / 、つまり \ / と折り返した場面なので、その時点のみ それまで続いていたものが「く」になっているか否かを判定すれば良い、ということになる(そして両者のカウントをリセットする)。ただし最終行が \ で終わった場合にその判定ができないので、最後に print するときに最終的な a==b を足してあげることで漏れをなくしている。

2022年11月 始まり (116 Byte)

yaneurao さんによる記録更新。

https://atcoder.jp/contests/iroha2019-day3/submissions/36137106 (116 Byte by yaneurao)

N,*C=open(0)
S=[0]*9**6
i=l=0
for c,*_ in C:i+=l!=c;S[i]+=(c>";")*2-1;l=c
print(sum(-a==b>0for a,b in zip(S,S[1:])))

解読

S にはそれぞれ / と \ の連続出現回数を交互に正負逆で入れていき、あとでその隣り合う数値を確認して同値であるものを数え上げる、という方式。
詳細は割愛。

2桁台に突入 (94 Byte)

yaneurao さんが自ら記録更新。

https://atcoder.jp/contests/iroha2019-day3/submissions/36192099 (94 Byte by yaneurao)

n,*A=open(0)
a=b=s=l=0
for c,*_ in A+[n]:
 if l!=c:s+=a==b>0!=";"<l;a=b;b=0
 b+=1;l=c
print(s)

解読

for c,*_ in とすることで改行を除く文字を確実に取得できる。
どうやらテストケースによってファイル末尾に改行があったりなかったりするようで、 open(0) から行ごとに取っていると "/\n" と "/" で同じはずなのに改行有無の違いで正しく判定できないことがある。

l に直前の c の値を入れながら for 文の中で比較していくことで、 / と \ が切り替わったタイミングを判定。
ただそれだけだと / → \ の変化か \ → / の変化かを判定できないので、";" より大きかったか否か、を見る。ASCIIで "/" と "\" の間であれば何でも良さそうではある。

a と b はそれぞれ 直前の連続出現回数と現在の連続出現回数を保持するもののようだ。

\ \ のように \ の連続で終わる場合にも判定をする必要があるので A+[n] として1回ぶん多くループを回す。
n は入力数を示す最初の数値行でしかないので必ず if 文の判定には入ることになる。なるほど…

正規表現で挑戦 (109 Byte)

このあたりのタイミングで私が参戦。
何かまったく違うアプローチで一発逆転できないかな、と思い、正規表現を利用した解法を試してみた。

https://atcoder.jp/contests/iroha2019-day3/submissions/36191733 (109 Byte by sugyan)

import re;print(sum([len(m[0])==len(m[1])for m in re.findall(r'(/+)(\\+)',"".join(map(str.strip,open(0))))]))

1行にはおさめることができたが あまり短くはならず、109 Byte…。

解説

まずすべての入力を繋げた "/\//\" のような文字列を作ってしまう。
これは "".join(map(str.strip,open(0))) で改行取り除いて繋ぎ直すことで得られる。

この文字列に対して re.findall(r'(/+)(\\+), ...)' とすることで 「/ が1回以上連続した後に \ が1回以上連続したもの」の塊を取得できる。
ので、それぞれの長さが等しいものだけを数え上げれば良い。

a{m}b{m} のように「同じ回数の繰り返し」を指定できる正規表現があれば長さの判定が要らなくて済みそうなのに…と思ったが、PCREにはあるがPythonの正規表現では使えないようだった。

https://regex101.com/r/CCEssq/1

正規表現 短縮版 (103 Byte)

どうにかもっと短くしたいと思い試行錯誤して、前述のを少しだけ縮めることができた。

https://atcoder.jp/contests/iroha2019-day3/submissions/36193754 (103 Byte by sugyan)

import re;print(sum([sum(map(ord,m))%139<1for m in re.findall(r'\\+/+',open(0).read().strip()[::-2])]))

解説

まず全部繋げた文字列を作るのに、 open(0).read().strip()[::-2] というものを使った。

"".join(map(str.strip,open(0))) だと 31 byte つかってしまう。他に open(0).read().replace("\n","") と read() して全文を取得してから改行だけ取り除く書き方もあるが、これも同じ文字数。re を使っている場合は replace の代わりに re.sub を使うことで 1 byte は短縮できるが、それでも微妙。

そもそも2行目以降はすべて / か \ のどちらか + 改行だけなので、1文字ずつとばして取得すれば目当てのものが取れるのではないか、と思い付いた。
ので open(0).read()[::2] のようなテクニックが使える。しかし1行目は何桁になるか分からず、ズレると困る。それなら逆に末尾から取ってしまえば良い。 [::-2] で末尾から1文字とばしで取得できる。
しかし今度はファイル末尾に改行があったりなかったりする問題にぶち当たるので、仕方ないので .strip() を入れてやる。それでも 28 byte と短い記述にすることができた。

逆から読むことになるので、正規表現の / と \ の順番も入れ替えてやる必要はある。

そして連続出現回数判定を、それぞれの長さを比較するのではなく sum(map(ord,m))%139<1 というズルい方法にしてみた。
/ \ それぞれの ord() を取ると 47 と 92 となる。

>>> ord("/") + ord("\\")
139

ので、/ と \ だけでできた文字列では「各文字の ord() を足し合わせたものが 139 で割り切れたとき」に同じ出現回数だと判定できそうだな、と。

考えれば分かるがこれは所謂「嘘解法」で、例えば / が 140回続いた後に \ が1回、のように出現回数の差が 139 の倍数だったときにも True としてしまうので正しい答えにならない。
が、試しにsubmitしてみたら通ってしまったので一応目的は達成している…。どうやらそういったテストケースは存在していなかったようだ。

余談

ちなみにこの sum(map(ord, ... )) % ... という ord の総和を使って判定するテクニックは自分のお気に入りで、以下の問題などでも同様の手法で最短記録更新に貢献した。(これは嘘解法ではない、はず)

https://atcoder.jp/contests/agc003/tasks/agc003_a

嘘解法による発展 (86 Byte)

正規表現の手法は最短記録は狙えなかったが、この ord() と % 139 を使った嘘解法、また入力を逆から見ていくといった手法を利用して徐々に最短記録が更新されていった。

https://atcoder.jp/contests/iroha2019-day3/submissions/36195164 (92 Byte by sugyan)

n,*A=open(0)
a=l=0;s=1
for c,*_ in A+[n]:
 if";">c!=l:a+=s%139<1;s=0
 s+=ord(c);l=c
print(a)

https://atcoder.jp/contests/iroha2019-day3/submissions/36195263 (91 Byte by yaneurao)

a=b=s=l=0
for c,*_ in[*open(0)][::-1]:
 if l!=c:s+=a==b>0!=";">l;a=b;b=0
 b+=1;l=c
print(s)

https://atcoder.jp/contests/iroha2019-day3/submissions/36195521 (89 Byte by yaneurao)

a=l=0;s=1
for c,*_ in[*open(0)][::-1]:
 if"/"<c!=l:a+=s%139<1;s=0
 s+=ord(c);l=c
print(a)

https://atcoder.jp/contests/iroha2019-day3/submissions/36195637 (88 Byte by sugyan)

a=l=0;s=1
for c,*_ in[*open(0)][::-1]:
 if"/"<c!=l:a+=s%139<1;s=0
 s+=ord(l:=c)
print(a)

https://atcoder.jp/contests/iroha2019-day3/submissions/36195891 (86 Byte by yaneurao)

a=l=s=d=0
for c,*_ in open(0):
 if";">c!=l:a+=d;s=0
 s+=ord(l:=c);d=s%139<1
print(a+d)

解読・解説

以前は / と \ の出現回数をそれぞれカウントしていって比較していたが、嘘解法を使うと単純に s に ord(c) を足し合わせていって / に折り返したときに 139 で割り切れるか否かだけを見る (そして 0 にリセットする)だけで良くなる。

また、最後の判定のためにループを1回多く回してやるところは [*open(0)][::-1] と逆から見ていくことで1行目の数値行を使うことで代用できるようになった。

しかしそれも「別の変数に代入しておいて最後にも足し合わせる」という形に変えることで逆順に見る必要もなくなり、かなり短くなった。

あとは s+=ord(c);l=c と計算後に代入するところを Python3.8 から使えるようになった Assignment expressions を使うことで s+=ord(l:=c) と書けて 1 byte だけ短縮できる。

if文の排除 (85 Byte)

ここまでの手法は for 文の中で if 文の分岐があるために1行では記述できず、どうしても改行後にインデントする必要があった。
これを排除するためにもう一つ変数を用意して bool 値と掛け算を使って計算するようにすることで 1 byte 短縮できた。

https://atcoder.jp/contests/iroha2019-day3/submissions/36291433 (85 Byte by sugyan)

a=l=s=d=0
for c,*_ in open(0):t=";">c!=l;a+=t*d;s-=t*s-ord(l:=c);d=s%139<1
print(a+d)

解説

それまで存在していた if";">c!=l:a+=d;s=0 という部分は

if (";" > c != l):
    a += d
    s = 0

ということで、「条件式を満たしたときだけ a に d を加えて s を 0 にする」という操作でしかない。ので、条件式の評価結果を b に代入すると

b = (";" > c != l)
a += d * b
s -= s * b

と書き換えることができる。 b が True のときだけ 2〜3行目の更新が実行されて、 False であれば何も変化が起きない。
この形への書き換えを利用することで、if 文を排除することができ、for 文を1行だけで記述することができる。

大きなブレークスルー (80 Byte)

ここで突如、嘘解法を使わない新しい手法で一気に記録が更新された。

https://atcoder.jp/contests/iroha2019-day3/submissions/36331780 (80 Byte by yaneurao)

a=l=s=0
for c,*_ in open(0):s-=[2,(l!=c)*~-s-2][c<"1"];a+=(s*s<2)*s;l=c
print(a)

解読

読みやすいように整形すると、

a = l = s = 0
for c, *_ in open(0):
    s -= [2, (l != c) * ~-s - 2][c < "1"]
    a += (s * s < 2) * s
    l = c
print(a)

~-s は正負逆転してビット反転、なので (s - 1) と同じになる、はず。
配列から c < "1" を index にして取っているのは分岐を短く書くためで、意味としては

    if c < "1":
        s -= (l != c) * (s - 1) - 2
    else:
        s -= 2

ということ。

  • / が最初に来たときには s -= 1 * (s - 1) - 2 なので s は 3 にセットされる
  • / が連続すると s -= 0 * (s - 1) - 2 なので s は 2 ずつ増加する
  • \ が来ると s は 2 ずつ減少する

…という形で、 s の値が 3 を起点に +2 もしくは -2 されて増減することになる。
で、 / と \ の数が釣り合ったときに 1 になり、それよりも \ が伸びていくと -1, -3, ... とどんどん負値になっていく。

出力すべきカウントは、 / の連続出現回数と \ の連続出現回数が釣り合ったときに +1 したいが、その時点では \ がまだ続くかどうか分からない。
これまでの手法では「再び新しい / が始まったとき」もしくは「最終行まで読み終わったとき」に判定して +1 していたが、

  • そこまでの / の連続出現回数と \ の連続出現回数が釣り合ったら とりあえず +1
  • その後 \ がもっと伸びたら -1 で取り消す
  • (新しく / が始まったら出現回数をリセットする)

という考え方にすれば、その都度処理できて「最終行まで読み終わったとき」など特別に考慮する必要がなくなる。

そして上述のように +1 すべきときと -1 すべきとき がそれぞれその値になっていれば、 a += (s * s < 2) * s で答えを更新していける。
これは絶対値が 1 以下のときだけその値を加える、ことになるのでそれ以外の場合は何も影響を受けない。

もし、s を単純に 0 起点で +1/-1 で増減させていると

if s == 0:
    a += 1
elif s == -1:
    a -= 1
else:
    pass

のような分岐を表現する必要があるが、2 ずつ変化させることでこのようにダイレクトに値をつかえるようになって短く書ける、というわけだ。

なるほど… いやぁこれは全然思いつかなかった…。

そこから符号逆転、簡略化 (76 Byte)

https://atcoder.jp/contests/iroha2019-day3/submissions/36348969 (79 Byte by sugyan)

a=l=s=0
for c,*_ in open(0):s+=[2,(l!=c)*~s-2][c<"1"];a-=(s*s<2)*s;l=c
print(a)

https://atcoder.jp/contests/iroha2019-day3/submissions/36359905 (76 Byte by sugyan)

a=l=s=0
for c in open(0):s+=[2,(l!=c)*~s-2][c<"1"];a-=(s*s<2)*s;l=c
print(a)

解説

前述の解法に対し、 s の符号と増減方向を逆転させると ~-s と書いていた部分が ~s だけにすることができる。
符号は逆になるが今度は a += としていたところを a -= に変えれば辻褄は合うので、これで 1 byte 縮めることができた。

意外に気付かなかったが、 for c,*_ in open(0): と受けていたところはもはや1文字目だけを取り出す必要は無くなっていた。
嘘解法のときのように ord() を使う場合や、明確に最後まで比較の条件として使われるなら必要だが、この解法の場合は l!=c を使うのは c<"1" のときだけで、つまり比較の誤判定が起きるとしても / で終わるときだけであり、どちらであったとしても s の値は a を変更し得るものではないので影響が無い。

まとめ

それまで 134 Byte が最短だったのが、このように様々な発展を遂げて数日間で 76 Byte まで縮められた。

執筆時点では Python3 での解答はこの 76 Byte が最短記録になっている。
まだもっと短く書けるかどうかは分からないが、この記事を読んだ方は是非記録更新に挑戦してみていただきたい(そして教えていただきたい…！)。

ちなみに言語問わずでの最短だと Bash で uniq と perl を使うもの。

https://atcoder.jp/contests/iroha2019-day3/submissions/5228396 (36 Byte by kotatsugame)

uniq -c|perl -pe'/\//|$`-$_||$\++}{'

(これもとんでもない解法だけど詳細は割愛。。)

Discussion