アルゴリズム問題に効くPython小ネタ

今週末の Google Code Jam Japan 2011 の決勝に向けて練習でもしておきたい所ですが、時間がなかなか取れないので、代わりにテクニックを紹介するブログエントリを書くことにしました。(代わりになってない、という意見は却下です)

namedtuple

namedtuple は python の標準ライブラリとして同梱されているクラスの一つです。

from collections import namedtuple

名前付きタプルという名が表す通り、タプルなのですが、タプルの各要素を名前で参照出来るようになります。たとえば、

Point3D = namedtuple('Point3D', 'x y z')
Range = namedtuple('Range', 'start end')

こんな感じで型を宣言します。そうすると、

p = Point3D(1, 2, 3)
r = Range(start=10, end=30)

こんな風に実体化できます。
もちろん、タプルなので、要素には添字アクセスできますし、それに加えて名前でもアクセスできます。

a = p[0] + p[1] + p[2]
a = p.x + p.y + p.z

もちろんタプルは immutable なので辞書のキーにしたり、setで多重値を取り除いたり出来ます。

d = dict()
d[p] = a

s = set()
s.add(Point3D(1, 2, 3))
s.add(Point3D(1, 2, 3))
assert(len(s) == 1)

ちょっと難しい話になりますので詳細は省略しますが、この型のオブジェクトは通常のオブジェクトよりも軽量に作られているので、大量のデータがあっても素の tuple に引けをとらない効率を誇ります。

使いどころとしては、専用のメソッドを持たなくても良いような簡単な構造データを表現するのに丁度良いかんじでしょうか。

Code Jam では、問題の入力データが構造を持っている時に、それを読み込む先を個別の配列にせず、コレで構造化しておくと、コードがすっきりして読みやすくなるでしょう。
N[i], M[i] のかわりに、a = r[i]; r.begin, r.end などとするわけです。

任意の条件でソート

さて、その namedtuple で定義したデータの配列などを持っている時、それのソートはどうすれば良いでしょうか。

単に tuple のリストをソート系関数に渡すと、タプルの各要素を順に比較し、最初に差が出た要素の大小で比較した結果でソートされます。

>>> range_list = [(i % 3, i, i % 2) for i in range(5)]
>>> range_list
[(0, 0, 0), (1, 1, 1), (2, 2, 0), (0, 3, 1), (1, 4, 0)]
>>> sorted(range_list)
[(0, 0, 0), (0, 3, 1), (1, 1, 1), (1, 4, 0), (2, 2, 0)]

この動作を調整するには、配列の sort() や sorted() に cmp や key といった引数を渡すことで操作できます。特に、key の方は上記の動作が頭に入っているとすごく判りやすく書けます。

たとえば、タプルの3番目の要素でソート(同じ値なら現在の順番でそのまま)なら

>>> range_list
[(0, 0, 0), (1, 1, 1), (2, 2, 0), (0, 3, 1), (1, 4, 0)]
>>> sorted(range_list, key=lambda x: x[2])
[(0, 0, 0), (2, 2, 0), (1, 4, 0), (1, 1, 1), (0, 3, 1)]

タプルの3番目の要素でソートし、同じ値なら1番目の要素、さらに同じなら2番目でソートするなら、

>>> range_list
[(0, 0, 0), (1, 1, 1), (2, 2, 0), (0, 3, 1), (1, 4, 0)]
>>> sorted(range_list, key=lambda x: (x[2], x[0], x[1]))
[(0, 0, 0), (1, 4, 0), (2, 2, 0), (0, 3, 1), (1, 1, 1)]

3つの要素の積でソートするなら

>>> range_list
[(0, 0, 0), (1, 1, 1), (2, 2, 0), (0, 3, 1), (1, 4, 0)]
>>> sorted(range_list, key=lambda x: x[0] * x[1] * x[2])
[(0, 0, 0), (2, 2, 0), (0, 3, 1), (1, 4, 0), (1, 1, 1)]

と、ラムダ関数と組み合わせて、最短の記述で自在なソートが実現できます。

ソートはアルゴリズム問題の基本なので、覚えておいて損はありません。

とはいえ、競技プログラミングでは最小さえ判れば良くて heapq が使われたりするので、意外と活躍の場は少ないんですが。(私のように富豪プログラミング寄りなコードを書く人間は高々100個程度な要素なら毎回ソートも気にしません。)