indexの配列を用いてpython辞書型やpandas.Dataframeから値を取得する

はじめに

やりたいことはタイトルの通り．
indexの配列があり，辞書やテーブルを参照して同じサイズのvalue配列を獲得したい．

やりかたはいろいろ

同様の課題がstackoverflowに存在した．そのうち，stackoverflowの回答で，リスト内包以外にもoperator，map関数などが比較されている．
indexの配列を直接代入するような短いコードでできそうかなと思ったが，そういった簡単な方法は見つからなかった．まあ回答の方法でスピードや可読性は十分ではある．

本記事では，リスト内包を用いた方法を中心にまとめる．

辞書型

indexの配列は2次元とした．
Jupyter notebookにて，%timeitで時間計測を行った．

import numpy as np

# 1-1000までのindexがランダムに入った1000 x 1000の配列
inds = np.random.randint(1,1001, (1000,1000))
# indexからvalueを得るための辞書
test_dict = { x: np.random.rand() for x in range(1,1001) }

# 手法1：一度numpyを平らにして一重リスト内包
%timeit np.array([test_dict[_] for _ in inds.flatten() ]).reshape(1000,1000)
# 手法2：numpy配列のまま二重リスト内包
%timeit np.array([test_dict[_] for __ in inds for _ in __]).reshape(1000,1000)
# 手法3：numpy配列をリストに変換してから二重リスト内包
%timeit np.array([test_dict[_] for __ in inds.tolist() for _ in __]).reshape(1000,1000)
'''結果：
514 ms ± 92.4 ms per loop (手法1)
468 ms ± 51   ms per loop (手法2)
229 ms ± 19.1 ms per loop (手法3)
'''

いずれのやり方もリスト内包で1次元の出力を得てから，reshapeしてもとの配列と同じ形にしている．tolist()を用いてリストに変換してからのほうが2倍以上速度が向上した（同様の事例は調べるといろいろ出てくる）．

pandas

pandasでは，値を取得するためにいくつかのメソッドのが存在する．
こちらもstackoverflowなどで速度について議論がなされており，この回答ではメソッドごとの速度比較を行っている．
早い順に，iat>at>iloc>locとなっているよう．

本記事でも，このうちいくつかの比較を行ってみた．

import pandas as pd

# 1-1000までのindexがランダムに入った1000 x 1000の配列
inds = np.random.randint(1,1001, (1000,1000))
# indexからvalueを得るためのpandasテーブル
test_df = pd.DataFrame({ 'INDEX': range(1,1001), 'VALUE': np.random.rand(1000) }).set_index(['INDEX'])

# 手法1：locを使用する
%timeit np.array([test_df.loc[_,'VALUE'] for __ in inds for _ in __]).reshape(1000,1000)
# 手法2：atを使用する
%timeit np.array([test_df.at[_,'VALUE'] for __ in inds for _ in __]).reshape(1000,1000)
# 手法3：iatを使用する
ind_dict = { _:i  for i,_ in enumerate(test_df.index) }
%timeit np.array([test_df.iat[ind_dict[_],0] for __ in inds for _ in __]).reshape(1000,1000)
'''結果：
13.1 s ± 211 ms per loop (手法1)
9.95 s ± 604 ms per loop (手法2)
12.3 s ± 290 ms per loop (手法3)
'''

うーん，辞書と比べてしまうとちょっと遅い．．．
この例ではatを用いる方法が最も早かった．iatを用いるために，行番号とindexの値とを結びつけるという強引なことをやってみたが，大幅な速度改善とまでは至らなかった．

おわりに

辞書型のほうが相当早いので，pandasテーブルを所得後，辞書型に変換して処理を行うのがよさそう．

new_dict = test_df['VALUE'].to_dict()

これ以上は並列計算とかCythonとかに譲るしかないかもしれない．