競技プログラミング用Python最速マスター

LLを使ってプログラミングするときに普段はRubyを使っているのですが,いろいろあってPythonも使い始めました。

何問かAOJでも提出してみたので、手元のメモ書きをリファレンス的に残しておこうと思います。

基本演算

-3 / 2 #=> -2
-4 % 3 #=> 2

入力

  • raw_input() で1行読み込み
  • string#split( [separator] ) で分割できる。半角スペースで分割したいときは、separatorを明示的に渡さずに 'hoge fuga piyo'.split() としておけばよい
  • "1 2 3"のような整数列を受け取りたいときは、次の2つが使える
map(int, raw_input().split())
[int(s) for s in raw_input().split()]

出力

Python 3系だとprint()関数Python 2系ではprintです。

print x # cout << x << endl;
print x, y # cout << x << " " << y << endl;
print x, # cout << x;

スコープ

割とややこしい。必要に応じてglobal宣言が必要なことがあります(自分で調べて)

リスト

l = [3, 1, 4]
l.append( 1 )
l.extend( [5, 9, 2] )

l #=> [3, 1, 4, 1, 5, 9, 2]

# リストの長さを取得
len(l) #=> 7

# スライス
l[2:5] #=> [4, 1, 5] (2から5-1まで)

# ソート
l.sort()
l #=> [1, 1, 2, 3, 4, 5, 9]

l.sort(reverse=True)
l #=> [9, 5, 4, 3, 2, 1, 1]

# スタックとして使いたい時に
l.pop() #=> 1
# キューとして使いたいときは,collections.dequeを使ったほうが良いです
l.pop(0) #=> 9

l #=> [5, 4, 3, 2, 1]

del l[1:4] # インデックス or スライスで指定する
l #=> [5, 1]

sort()の引数には、次のようなものがオプションとして付けられます

引数の名前 説明
cmp 比較関数。値を2つ取ってboolを返す関数 or lambda式を渡す
key 値を1つ取って、値を変換する関数 or lambda式を渡すと、それを基にソートしてくれる (SQLでいうORDER BY)
reverse Trueを指定すると、逆順ソートになる

高階関数

  • filter
  • map
  • reduce

などなど

filter(lambda x: x % 2 == 0, xrange(10))
#=> [0, 2, 4, 6, 8]

リスト内包表記

[str(x) for x in xrange(10) if x % 2 == 0]
#=> ['0', '2', '4', '6', '8']

リストの罠

リストの生成

参照がコピーされる場合があるので、注意しましょう!

# 1次元の場合(配列の中身がプリミティブ型の時?)
memo = [0] * 100

# 多次元の場合
# これはダメ(参照がコピーされる)
memo = [[0] * 10] * 10
# こっちにする
memo = [[0] * 10 for x in xrange(10)]

リストのコピー

参照のコピーではなく、値のコピーを行いたい時には注意!

# 1次元の場合
v = [1, 2, 3, 4, 5]

# これはダメ(参照がコピーされる)
w = v
# 値をコピーしたいなら、こっち
w = v[:]

# 多次元の場合
v = [[1, 2, 3], [4, 5, 6]]

# これはダメ(内部のリストは、結局参照でコピーされる)
w = v[:]
# deep copyを使う
from copy import deepcopy
w = deepcopy(v)

束縛,浅いコピー (shallow copy),深いコピー (deep copy) については,Rubyの説明ですが以下のページがおすすめです。

qiita.com

その他データ構造

set

set関数にコレクションを渡してやると、集合を表すコレクションを作ることができます。 -&|^ などの演算子が定義されているので、困ることはないでしょう。

dictionary

波括弧{} に key:value の形で入れることで作ることができますが、keyとして使えるのは数値、文字列、「数値か文字列によって構成されるタプル」などの不変オブジェクトのみ。

キーの一覧はkeys() メソッドで取得できる。

そのキーが含まれているか判定するときは has_key() メソッドか、in演算子で判定する。

dict関数にタプルのリストを渡したり、キーワード引数を用いて dict(foo=1, bar=2) のようにして生成することもできる(ただし、keyが文字列の時のみ)

key:value を同時に取り出したいときは、iteritems() メソッドを使う。

for k, v in d.iteritems():
    print k, v
存在しない要素にアクセスする

C++でワードカウントのような操作を行うときに、次のようなコードを書いている人も多いと思います。

std::map<std::string, int> word_count;

word_count[word]++;

Pythonではディクショナリに存在しない要素を読み出そうとした時に、KeyError という例外が出ます。

has_key() や in演算子を使わずに書く方法としては、次のような方法があります

collections.Counterクラスを使う (おすすめだけど、存在をよく忘れる)

Counterクラスは、デフォルトのvalueが0のディクショナリです。

from collections import Counter

c = Counter()

c['hoge'] += 1
c['fuga'] += 1
dict#setdefault() を使う (何も考えたくないときにおすすめ)

setdefault(key, value) はkeyが存在しないときはvalueで初期化するメソッドです

c = {}

c.setdefault('hoge', 0) # c['hoge'] を0で初期化
c['hoge'] += 1

c.setdefault('hoge', 0) # 既にc['hoge']が存在するので、何もしない
c['hoge'] += 1
collectioins.defaultdict を使う

リファレンスへのリンク

defaultdictはdictのサブクラスで、存在しないkeyの値を読もうとしたり書き換えたりしようとした時に、コンストラクタとして与えた関数や無名関数を呼び出して、その返り値をデフォルト値とします。

from collections import defaultdict

s = "jan feb mar apr may"
month = defaultdict(lambda: len(month))
for w in s.split():
  month[w]

month # => defaultdict(<function <lambda> at 0x1039666e0>, {'jan': 0, 'apr': 3, 'mar': 2, 'feb': 1, 'may': 4})

Queue & Stack & Priority Queue

この3つのコレクションは、2種類あるようです

  1. collections.dequeheapq を使う方法

  2. queue モジュールに含まれる Queue, LifoQueue, PriorityQueue を使う方法

前者のほうがデバッグが楽そうなので (repr() に渡すだけで中身が見れる) 、前者を使うことにします。

追記

queueモジュールのコレクションは、本来はマルチスレッドでのプログラミングのために用意されたものだと書いてあったので、シングルスレッド環境でベンチマークを取ってみました。

その結果、collections.dequeやheapqと比べてかなり重い事がわかりました(ロックと解放に時間がかかっているんでしょう)。手元の環境 (Core i5 1.7GHz, 4GB, Mac OS X Lion, Python 2.7.1) で行った即席のテストでは、heapq と PriorityQueueでは2倍ほど、dequeとQueue, LifoQueueについては5倍ほどの差が出ました。

競技プログラミングにおいては、queueモジュールを選択するメリットが少ないように思われます。

upper / lower bound

bisect という便利なモジュールがあります

省メモリ & 高速化のために…

array

リストに整数を入れたときは、1要素あたり16バイト(!)確保されるので、たまにMLEすることがあります。

そういうときは、省メモリなarrayモジュールが使えます。ただ残念なことに、sort() などの必須とも言えるいくつかのメソッドが定義されていません。

from array import array
a = array('i', [3, 1, 4, 1, 5])

1番目の引数には、格納する型を表す文字列を渡す必要があります。 詳細はこちらで確認して下さい。

8.6. array — Efficient arrays of numeric values — Python v2.7.3 documentation

覚えておくのは 'b', 'i', 'l', 'd' くらいで良いんじゃないかなー。大文字にすると unsigned になることも覚えておけば、困ることは無いかと。

ジェネレータ式

リストを引数に取る系の関数は、リストの実態の代わりに、ジェネレータ式というものを渡すことができます。

# リストの実体の中間生成データができるので、メモリが厳しい…
a = array('i', [ 0 for x in xrange(1000) ] )

# []の代わりに()でくくりましょう!
a = array('i', ( 0 for x in xrange(1000) ) )

countif

リストlの中でnより大きい要素の個数を数える時

ダサい方法ですが、これが一番省メモリでした。

n = 5
l = [2, 7, 1, 8, 2, 8]

c = sum(1 for x in l if x > n)