読者です 読者をやめる 読者になる 読者になる

渡米生活。(日記)

渡米生活。本家から切り離しました。あまり渡米生活に関係のないプログラムネタや音楽ネタなど。

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というのだそうで。詳しくは公式ページを見て下さい。

8.6. 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個あるとしたら、0番目から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関数は、デフォルトでは各要素の0番目の値で
#ソートしてくれる。np.arrayでは違った振る舞いをするので注意。

# wmapからstd::mapのキーに相当する各要素の0番目の値だけを抜き出した配列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]