どうもこんにちはヘボプログラマです。
なんとなしに競プロのAtCoderというサイトに登録していたのですが、
ある日メールが来て、「今日の21時からコンテストやるよ!!」と言われたので、鼻息荒く参加してみました。
今回参加したのはこれです。
AtCoder Grand Contest 023
リアルタイムで参加するのは初めてだったので、どれほどのものか腕試し程度に頑張ってみようと思いました。
問題はAからFまであり、Aが最も簡単であとに行くほど難しくなり得点が高くなります。
僕は普段はPythonとかJSとかを主に書くのですが、JSはフレームワークの使い方を知っている程度で
言語に対しての理解は微妙なのでPythonで臨んでみました。
A(200)とC(800)を解いて1000点を狙ったのですが、結果は・・、
A
C
解けてるのに!!TLE!!
TLEというのは「Time Limit Exceeded」の略で、要するに「時間かかり過ぎだよアホ乙」という意味です。
こうなると結局得点はもらえず、0点になります。
結局2時間20分の時間制限の中奮闘しましたが、0点で終わりました。
ということで大事なのはここからです。
AtCoderは人が提出したコードを見ることができます。
それを教材にさせてもらって、僕のなんの面白みのないクソコードをどう改善すればいいのか、
どんなアルゴリズムの組み方が考えられるのか、
と言ったところを勉強していきたいと思います。
問題
問題文
長さ N の整数列 A があります。A の 空でない 連続する 部分列であって、その総和が 0 になるものの個数を求めてください。 ただし、ここで数えるのは 部分列の取り出し方 であることに注意してください。 つまり、ある 2 つの部分列が列として同じでも、取り出した位置が異なるならば、それらは別々に数えるものとします。
考え方
公式から考え方の例が掲載されているので参考にします。
【参考】
AtCoder Grand Contest 023 解説放送 – YouTube
https://img.atcoder.jp/agc023/editorial.pdf
どうやら今回の問題を通して学ぶべきキーワードは
- 計算量オーダーという概念
- 累積和
の2点っぽいです。
計算量オーダーについて
アルゴリズム界隈でよく見るO(n)といった記法のことです。
アルゴリズムの計算量をざっくり測る指標です。
詳しくは以下のサイトを参考にしてください。
【参考】
オーダー初見の時に見るとわかりやすい
[初心者向け] プログラムの計算量を求める方法 – Qiita
オーダー知った後に思い出す時に役立ちそう
計算量オーダーについて – Qiita
では、僕が提出したコードがどれほど効率が悪いのか見てみます。
1 2 3 4 5 6 7 8 9 10 |
n = int(input()) s = list(map(int, input().rstrip().split())) count = 0 for i in range(n): for k in range(i+1,n+1): if sum(s[i:k]) == 0: count += 1 print(count) |
二重forループのオーダーが
なので
O(n^2)とかになるのでしょうか。
if文の中身はオーダーの計算の仕方がわからない・・。
調べてるとこんな記事がありました。
Pythonistaなら知らないと恥ずかしい計算量のはなし – Qiita
ここを見てみるとsliceはO(k)らしいです。
~中略~
というわけでO(n^3)になるっぽいです。(たぶん)
確かにTLEになりそうですね。
累積和
次に累積和という考え方を見ていきます。
以下のサイトが分かりやすかったです。
【参考】小学生でもわかる累積和と中学生でもわかるBinary Indexed Tree – anctgcc’s diary
要するに、部分列の和や差を求めるときには、逐次計算をするのではなく
累積和をとったリストを別に用意することで、計算回数を削減できるというテクニックです。
これを頭に入れた上で強い人のコードを見てみます。
Python
強い人のコードその1
まずはhiro_hiroさんのコードです。
めっちゃシンプルです・・。すごい・・。
入力の部分を除くと実質2行ですね・・。
1 2 3 4 5 6 7 8 |
from itertools import accumulate from collections import Counter n = int(input()) a = [0] + list(map(int, input().split())) s = list(accumulate(a)) print(sum([si * (si - 1) // 2 for si in Counter(s).values()])) |
ちょっと詳細を見てみます。
Pythonの標準モジュールに累積和を求めるaccmulate()というのがあるようです。
1 2 3 4 5 |
>>> from itertools import accumulate >>> a = [0, 1, 3, -4, 2, 2, -2] >>> s = list(accumulate(a)) >>> s [0, 1, 4, 0, 2, 4, 2] # 累積和のリスト! |
またこちらも標準モジュールにCounter()というのがあり、
リストの中にある要素がいくつずつあるのかを出力してくれます。
1 2 3 4 5 |
>>> from collections import Counter >>> Counter(s) Counter({0: 2, 4: 2, 2: 2, 1: 1}) >>> Counter(s).values() dict_values([2, 1, 2, 2]) |
ここまで来るとゴールは目前で、[2, 1, 2, 2]の各要素に対して、
Combination 2を求めます。
は、ループなどを回さずとも、
で求めることができますね。
コレの和が解になります。
これを一行で表記すると
1 |
print(sum([si * (si - 1) // 2 for si in Counter(s).values()])) |
となります。
あと2つほど強い人のコードを見てみます。
強い人のコードその2
次はへーさんさんのコードです。
前述した記事に書いてありましたが、dict型を用いると、ハッシュテーブルを使用することになるので、高速に参照できるようです。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
n = int(input()) a = list(map(int, input().split())) ans = 0 l = {0:1} s = 0 for i in a: s += i if s in l: ans += l[s] l[s] += 1 else: l[s] = 1 print(ans) |
8行目で累積和を作っていますね。
そのままその値が初見なら、その値をキーにして辞書に登録。
その値が既知なら辞書の値を足していっています。
ここで、おお!!すごい!!って思ったのは、
の値って、
なので、
と同値になります。
つまり、1+2+…+(n-1)になるんですよね。
それをさらっと使っています。勉強になります・・。
強い人のコードその3
最後にtjakeさんのコードを見てみます。
1 2 3 4 5 6 7 8 9 10 11 |
N = int(input()) *A, = map(int, input().split()) B = {} s = 0 ans = 0 B[0] = 1 for a in A: s += a ans += B.get(s, 0) B[s] = B.get(s, 0) + 1 print(ans) |
僕初めて見たのですが、辞書型をいいカンジに操作するgetメソッドというのがあるのですね。
【参考】[Python] getメソッドを使って辞書の値を取得する | アシベパンチ
これを用いることで、先ほどのコードのif文をより簡潔に書かれています。
Nim
せっかくなのでNimでも書いてみましょう。
AtCoder内にNimユーザーがあまりいないのと、そもそも僕がNimのことを全然知らないので、
Nimのイイカンジのモジュールはあまり使わず(というか使えず)、Pythonコードの書き換えのようになってしまいました。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import strutils, tables, sequtils var n = stdin.readLine.parseInt A = stdin.readLine.split.map(parseInt) l = initTable[int, int]() s = 0 ans = 0 l[0]=1 for a in A: s += a if s in l: ans += l[s] l[s] += 1 else: l[s] = 1 echo ans |
Nimでは辞書型のことをTableと呼ぶようです(?)
【参考】Nimのデータ型 – Nim帳
Nimで書いて提出してみたところPythonで書かれたコードの中で一番速いものと同じタイムでした。
あれ、。
今絶賛勉強中のHaskellでもやってみたかったのですが、
あまりにも知らないのでもう少し後でチャレンジしてみます・・。
以上。