Pythonで辞書型データから高速に検索したい

Question

以下の辞書型のデータ構造で，valueから検索したいケースが生じてしまい，良い検索手段を検討しております。
高速な方法があればご教授いただきたいです。

以下のような文字列をkeyとし，文字列が格納されたリストがvalueになった辞書型のデータ構造で，数万レコードあります。

なお，リスト[0]がkeyと同値になっています。

test_dict =
{"A":["A","B","C"],
 "B":["B","C","D"],
 :
 "X":["X","Y","Z"]}

これまではこの辞書型データに対して，keyを利用した検索を行っておりました。

>print(test_dict["A"])
["A","B","C"] 

しかし，リスト[1]が"B"のリストを取得したいという事態となりました。

これを純粋にやると以下のようになり性能懸念があり，改善策を検討しております。

index = 1     # 取得したい値のインデックス
target = "B"  # 取得したい値

for key, value in test_dict.items():
    if target == value[index]:
        print(value) # ["A","B","C"]

良い方法があればご教授ください。

Answer 1

検索回数によってやるべきことが変わります。

一度しか検索しないのであればO(N)を我慢するしかありません。

何度も検索するのであれば、最初にそのための辞書を作ってからの方が速くなります。

---

衝突回避版。

test_dict = {
    "A":["A","B","C"],
    "B":["B","C","D"],
    "X":["X","Y","Z"],
}

test_dict_1 = {}
for k, v in test_dict.items():
    if v[1] in test_dict_1:
        test_dict_1[v[1]].append(k)
    else:
        test_dict_1[v[1]] = [k]

print([test_dict[k] for k in test_dict_1['B']])

Answer 2

こうやってあらかじめデータ構造を変更してみればいかがでしょうか？

test_dict = {
    "A":["A","B","C"],
    "B":["B","C","D"],
    "X":["X","Y","Z"],
}
 
test_dict_1 = {v[1]: v for v in test_dict.values()}
print(test_dict_1)

Comments

[mkgrei]

keyが複数存在した場合が少し面倒ですね。

[YouheiSakurai]

そうですね。前提のデータによってはそのコーナーケースに殺される可能性はありますね。

[badtz]

ご回答ありがとうございます。
print(test_dict_1["B"])と理解しました。
test_dictのkeyはユニークで，リストの要素数はすべて同じになっています。
何度も検索するのですが，毎回検索したいインデックスと値が異なるので一度切りです。
list[1]の”B"，list[2]の”Z"，list[1]の”C"と検索したい時に，
list[1]の”B"，list[1]の”C"，list[2]の”Z"の順番に変更する(インデックス毎に分ける)と，上記の方法で早くなりそうです。

Answer 3

Pythonの機能だけだと改善のしようがないので、pandasのDataFrameを使った方法で回答します。

import pandas as pd
test_dict = {
     "A":["A","B","C"],
     "B":["B","C","D"],
     "X":["X","Y","Z"]
}
df = pd.DataFrame.from_dict(test_dict, orient="index")
print(df)

#    0  1  2
# A  A  B  C
# B  B  C  D
# X  X  Y  Z

これで今までのデータをDataFrameに移行しました。
ここから、

index = 1     # 取得したい値のインデックス
target = "B"  # 取得したい値

result = df[df[index] == target]
result = result.values.tolist()  # コメントアウトしても動きます
print(result)
# [['A', 'B', 'C']]

などと出来ます。結果が配列の配列になっているのは、該当するレコードが２個以上になる場合もあるためです。
ただ、kyeアクセスには少しだけ手間が増えますので、場合によってはリファクタリング箇所が増えてしまう恐れがあります。

key = "X"
value = df.loc[key].tolist()
print(value)
# ['X', 'Y', 'Z']

ご参考までに。

---

追記
速度を測ってみました。結論から言うと、1回の検索ですらpandasのほうが平均して遅いという結果になりました。

適当にデータを作成

import pandas as pd
from itertools import product
temp = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
test_dict = {s+t+u: [s+t+u+"0", s+t+u+"1", s+t+u+"2"] for s, t, u in product(temp, repeat=3)}
df = pd.DataFrame.from_dict(test_dict, orient="index")
print("head")
print(df.head())
print("tail")
print(df.tail())
print("レコード数", len(df))

# head
#         0     1     2
# aaa  aaa0  aaa1  aaa2
# aab  aab0  aab1  aab2
# aac  aac0  aac1  aac2
# aad  aad0  aad1  aad2
# aae  aae0  aae1  aae2
# tail
#         0     1     2
# ZZV  ZZV0  ZZV1  ZZV2
# ZZW  ZZW0  ZZW1  ZZW2
# ZZX  ZZX0  ZZX1  ZZX2
# ZZY  ZZY0  ZZY1  ZZY2
# ZZZ  ZZZ0  ZZZ1  ZZZ2
# レコード数 140608

これに対して元の方法とpandasを比較します。

def get_value0(index, target):
    # 元の方法
    for key, value in test_dict.items():
        if target == value[index]:
            return value

def get_value1(index, target):
    # pandas使った方法
    return df[df[index] == target].values.tolist()[0]

# 探したいデータ        
index = 2     # 取得したい値のインデックス
target = "Gcw2"  # 取得したい値

# 同じ結果になるか確認
assert get_value0(index, target) == get_value1(index, target)

# 元の方法
%timeit get_value0(index, target)
# 7.78 ms ± 131 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# pandasを使った方法
%timeit get_value1(index, target)
# 11.7 ms ± 231 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

targetによっては元の方法では最後までloopを回さないとダメなにで、パフォーマンスにムラがあります。ただ、ワーストケースでpandasより1ms遅い程度でした。

実行環境
Python 3.6.4
pandas==0.22.0
MBP

Comments

[badtz]

ご回答ありがとうございます。
こちらの方法は大変参考になりました。
行列のように取り扱うことで高速にアクセスできるということでしょうか。
少し調査して試してみたいと考えております。

[退会済みユーザー]

pandasの内部がPythonよりも最適化されているようですので、それに頼った方法です。どの程度差が出るかはまだ測っていません。

[mkgrei]

pandasの方が実装は速いのですが、比較演算子がN回実行されるので、同じデータセットに対して2度以上検索をかけると辞書を一度作った方が速くなります。

代わりにデータベースのように一括で実行できるようになるので、やりたいこと次第では工夫することで汎用性高いまま高速化できます。

[退会済みユーザー]

おっしゃる通りでした。試しに適当に約14万件のレコードを作って試しましたが、1回の検索ですらpandasが特に速いという結果にはなりませんでした。何か間違えているかもしれません。

[退会済みユーザー]

時間測った結果を追記しました。あんまり良い結果ではないので、この回答は参考程度にとどめておいてください。

[mkgrei]

焼け石に水ですが、
df[df[index].values == target].values.tolist()[0]
で少し速くなります。

[badtz]

測定もしていただきありがとうございます。しかし本方式のほうが遅くなってしまうのは残念です。
別の検索要件が出た際には汎用化できそうとのことですので、再検討させていただこうと思いますが、一時的な辞書を作成する方式で比較してみたいと考えております。