hatunina’s blog

メモと日記です

pandas.DataFrameからlistの要素を取得したい

あんまり使う機会ないけど毎回忘れるのでメモです。

こんな感じのitemsがリストになっている要素を取得したいやつです。

df.head()

#   id  items
# 0     A   [A, B]
# 1     B   [B, C, D]
# 2     C   [B, C]

print(df.iloc[2, 1])
print(type(df.iloc[2, 1]))

# ['B', 'C']
# <class 'list'>


とりあえず普通に比較して要素を取得してみると

val_list = ['B', 'C', 'D']
df[df['items']==val_list]

# 何も返ってこない

val_list = ['B', 'C']
df[df['items']==val_list]

# ValueError: Arrays were different lengths: 3 vs 2


val_list = ['B', 'C', 'D']の場合、データフレームの1行目からB, 2行目はC, 3行目はDというようにリストの要素と各行を比較してしまうのでマッチするものがなく何も返ってこない。
val_list = ['B', 'C']の場合、行数と要素数が合っていないのでエラーが出る。
ぼーっとしてると見逃しそう。。。


applyを使って1行ずつ比較することで解決!

val_list = ['B', 'C']
df[df['items'].apply(lambda x: x==val_list)]

# id    items
# 2     C   [B, C]