強い人のコードから学ぶアルゴリズム(AtCoderGrandContest-023-A)

どうもこんにちはヘボプログラマです。

なんとなしに競プロのAtCoderというサイトに登録していたのですが、
ある日メールが来て、「今日の21時からコンテストやるよ！！」と言われたので、鼻息荒く参加してみました。

今回参加したのはこれです。
AtCoder Grand Contest 023

リアルタイムで参加するのは初めてだったので、どれほどのものか腕試し程度に頑張ってみようと思いました。

問題はAからFまであり、Aが最も簡単であとに行くほど難しくなり得点が高くなります。
僕は普段はPythonとかJSとかを主に書くのですが、JSはフレームワークの使い方を知っている程度で
言語に対しての理解は微妙なのでPythonで臨んでみました。

A(200)とC(800)を解いて1000点を狙ったのですが、結果は・・、

解けてるのに！！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

では、僕が提出したコードがどれほど効率が悪いのか見てみます。

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)

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ループのオーダーが
$\sum_{k=1}^{n}k = \frac {1}{2}n(n-1)$ なので
O(n^2)とかになるのでしょうか。
if文の中身はオーダーの計算の仕方がわからない・・。

調べてるとこんな記事がありました。
Pythonistaなら知らないと恥ずかしい計算量のはなし – Qiita

ここを見てみるとsliceはO(k)らしいです。

~中略~

というわけでO(n^3)になるっぽいです。(たぶん)
確かにTLEになりそうですね。

累積和

次に累積和という考え方を見ていきます。

以下のサイトが分かりやすかったです。
【参考】小学生でもわかる累積和と中学生でもわかるBinary Indexed Tree – anctgcc’s diary

要するに、部分列の和や差を求めるときには、逐次計算をするのではなく
累積和をとったリストを別に用意することで、計算回数を削減できるというテクニックです。

これを頭に入れた上で強い人のコードを見てみます。

Python

強い人のコードその1

まずはhiro_hiroさんのコードです。
めっちゃシンプルです・・。すごい・・。

入力の部分を除くと実質2行ですね・・。

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()]))

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()というのがあるようです。

>>> from itertools import accumulate
>>> a = [0, 1, 3, -4, 2, 2, -2]
>>> s = list(accumulate(a))
>>> s
[0, 1, 4, 0, 2, 4, 2] # 累積和のリスト!

>>> 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()というのがあり、
リストの中にある要素がいくつずつあるのかを出力してくれます。

>>> from collections import Counter
>>> Counter(s)
Counter({0: 2, 4: 2, 2: 2, 1: 1})
>>> Counter(s).values()
dict_values([2, 1, 2, 2])

>>> 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を求めます。
${}_n C _2$ は、ループなどを回さずとも、
$\frac {1}{2} n(n-1)$ で求めることができますね。

コレの和が解になります。
これを一行で表記すると

print(sum([si * (si - 1) // 2 for si in Counter(s).values()]))

1	print(sum([si * (si - 1) // 2 for si in Counter(s).values()]))

となります。

あと2つほど強い人のコードを見てみます。

強い人のコードその2

次はへーさんさんのコードです。
前述した記事に書いてありましたが、dict型を用いると、ハッシュテーブルを使用することになるので、高速に参照できるようです。

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)

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行目で累積和を作っていますね。
そのままその値が初見なら、その値をキーにして辞書に登録。
その値が既知なら辞書の値を足していっています。

ここで、おお！！すごい！！って思ったのは、
${}_n C _2$ の値って、
$\frac {1}{2} n(n-1)$ なので、
$\sum_{k=1}^{n-1}k$ と同値になります。
つまり、1+2+…+(n-1)になるんですよね。
それをさらっと使っています。勉強になります・・。

強い人のコードその3

最後にtjakeさんのコードを見てみます。

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)

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コードの書き換えのようになってしまいました。

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

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でもやってみたかったのですが、
あまりにも知らないのでもう少し後でチャレンジしてみます・・。

以上。

強い人のコードから学ぶアルゴリズム(AtCoderGrandContest-023-A)

問題

考え方

計算量オーダーについて

累積和

Python

強い人のコードその1

強い人のコードその2

強い人のコードその3

Nim

関連

コメントを残すコメントをキャンセル

問題

考え方

計算量オーダーについて

累積和

Python

強い人のコードその1

強い人のコードその2

強い人のコードその3

Nim

共有:

関連

コメントを残す コメントをキャンセル

コメントを残すコメントをキャンセル