Pythonの配列でC++のstlのupper_boundやlower_bound みたいなことをやるには
2013/9/17 mapの場合を文末に追記。
久々にpython tips.
「python」「検索」「アルゴリズム」などのキーワードで検索すると、文字列検索の情報しかヒットしないのですが、ソートされたpythonの数値配列で、「ある値と同じか大きい」みたいな検索をしたいことがあります。
私は普段C++を使っているので、こういうときはstl::setで配列を作って、upper_bound()かlower_bound()を使えば一発なのですが、これをpythonでやるには多少変更が必要なようです。
pythonの二分法アルゴリズムモジュールは、bisectというのだそうで。詳しくは公式ページを見て下さい。
まあ、上のページを見ればそれでおわりなんですが、上に書いたC++での手法をそのままやろうとすると失敗します。
というのは、pythonにもsetがあるんですが、これはindexを保持しないので、bisectモジュールが使えません。
というわけで、同等のことをやろうとすると、以下の通りです。
import bisect #配列を用意。作りながら昇順にソート #(勿論、適当に配列を作ってソートしてもOKだけど、 #bisectの例題なので、insort_leftを使ってみた。） array = [] bisect.insort_left(array, 9) bisect.insort_left(array, 1) bisect.insort_left(array, 7) bisect.insort_left(array, 3) bisect.insort_left(array, 11) bisect.insort_left(array, 5) #配列を書き出してみるとソートされている print array [1, 3, 5, 7, 9, 11] # 検索 （戻り値はインデックス） bisect.bisect_left(array, 8) 4 bisect.bisect_left(array, 9) 4 bisect.bisect_right(array, 8) 4 bisect.bisect_right(array, 9) 5
この振る舞いをみるところ、lower_boundはbisect_leftに、upper_boundはbisect_rightにそれぞれ相当するようです。
追記：
さて、上記一次元配列は簡単なのだけど、C++のstd::mapのlower_boundに相当することをやろうとすると、ひと手間かかります。
Python公式ページによれば、以下のようにしろ、とのことです。
例題として、Monte Carloでよくやる重みつきのrandom pickupをやってみます。とりあえずrandom numberの配列を作って、それぞれの数に比例した割合でどれか一つを選び出すプログラム。
C++ではどうするかというと、まず重みデータが10個あるとしたら、０番目からi番目までの値を足したものをstd::mapのi番目のkeyに、そしてi番目の値をi番目のvalueにセットする。そうすると、サイズ10のmapができます。そうしておいて、ランダムナンバーを0からmapのkeyの最後の値（つまり重みデータ全部の和）の間でふり、得られたランダムナンバーがmapのkey配列のどこに相当するかを調べる、という感じです。
import numpy as np import random import bisect # ウェイトデータを準備 data = np.random.uniform(0, 10, 10) # ウェイトのsumを計算 wmap = [] wsum = 0 for w in data : wsum += w wmap.append([wsum, w]) print data print wmap #この例題の場合はwmapが既に各要素の0番目の値でソートされているが、 #もしソートが必要ならここでソートをかける #pythonのリストのsort関数は、デフォルトでは各要素の０番目の値で #ソートしてくれる。np.arrayでは違った振る舞いをするので注意。 # wmapからstd::mapのキーに相当する各要素の０番目の値だけを抜き出した配列keysを作る。 keys = [r[0] for r in wmap] print keys # random number取得 random = random.uniform(0, wsum) # bisectでrandomが入るべきindexを探す. wmapでなくkeysに対して実行 index = bisect.bisect_left(keys, random) print index # wmapの相当するレコードを探す print wmap[index]
