【Pandas】データフレームから特定の行を抽出し、新しいデータフレームに高速に移動[Python]

2023年2月9日2023年3月10日

URLをコピーしました！

Pandas

前回、リスト内の各要素の個数を数える方法（count、collections.Counter）を紹介しました。

あわせて読みたい

【Python基礎】リスト内包表記：二次元リスト（ネストされたリスト）の作成【リスト内包表記】前回、Twitterでのツイートの文字数を正確に数えてくれるライブラリtwitter-text-parserを試してみました。今回は話題がガラッと変わってリスト内…

今回はPandasでデータフレームから特定の行を抽出し、新しいデータフレームに高速に移す方法を紹介します。

なぜこれを書こうかと思ったかというと、これまではfor文で抽出し、それを1行1行新しいデータフレームに追加していっていました。

しかしこの方法は実はかなり遅いということが分かり、試行錯誤したところ、抽出した行の行番号をリストの格納し、その番号を持って新規のデータフレームを作成する方がはるかに速いということが分かりました。

そしてせっかくなのでこの方法を紹介しておこうと思ったのが、今回の話の発端です。

ということでまずは今回使用するデータフレームをこんな感じで作成しました。

import random
import pandas as pd

data = [[random.randint(0,10), random.choice([True, False])] for i in range(100)]
df = pd.DataFrame(data, columns=['val', 'bool'])

print(df)

実行結果
    val   bool
0     4  False
1     5   True
2     5  False
3     0  False
4     9  False
..  ...    ...
95    0  False
96   10  False
97    0  False
98    4  False
99    1  False

[100 rows x 2 columns]

数値とbool値（True/False）をもつデータフレームです。

それでは始めていきましょう。

Pandasの基本のデータ抽出方法

まずPandasの基本のデータ抽出方法です。

例えば先ほどのデータフレームでboolの値がTrueのものだけ抽出する場合、「df[抽出条件]」という形で抽出します。

df_true = df[df['bool'] == True]

print(df_true)

実行結果
   val  bool
0     0  True
2     9  True
7     7  True
（中略）
97    2  True
98    9  True
99    4  True

また複数条件を指定する場合は、それぞれの条件を括弧でくくり、「&」（AND）、「|」（OR）を用います。

df_true5 = df[(df['bool'] == True) & (df['val'] == 5)]

print(df_true5)

実行結果
    val  bool
8     5  True
15    5  True
37    5  True
41    5  True
52    5  True
71    5  True
80    5  True

for文を使った抽出方法

ただ条件が複雑になればなるほど、上記の方法は分かりにくくなります。

またセルにリストを使っていたりすると、さらに抽出が難しくなります。

そこでfor文を使って、1行ずつ取り出し、条件に合うかどうか確認し、合うものを新しいデータフレームに移すなんてことをすることもあるでしょう。

そんな時はこんな感じになるのではないでしょうか。

df_for1 = pd.DataFrame()
for i in range(len(df)):
    if df['bool'].iloc[i] == True:
        if df['val'].iloc[i] == 5:
            df_for1 = pd.concat([df_for1, df.iloc[i]], axis=1)
            
df_for1 = df_for1.transpose()
    
print(df_for1)

実行結果
   val  bool
8    5  True
15   5  True
37   5  True
41   5  True
52   5  True
71   5  True
80   5  True

1行ずつデータを確認し、条件に合うものを「concat」を使って、新しいデータフレームに結合していっています。

この方法は何度か使っていたのですが、なんとも遅い。

ということで試してみたのが、こちらの方法です。

df_rownum = []
for i in range(len(df)):
    if df['bool'].iloc[i] == True:
        if df['val'].iloc[i] == 5:
            df_rownum.append(i)
    
df_for2 = df.iloc[df_rownum]

print(df_for2)

実行結果
    val  bool
8     5  True
15    5  True
37    5  True
41    5  True
52    5  True
71    5  True
80    5  True

こちらの方法では条件に合う行の行番号をまず取得、リストに格納しています。

そしてその行番号のリストを使って「iloc」で一括に行を取得、新しいデータフレームに挿入をしています。

それぞれ「%%timeit」で処理速度を算出してみました。

%%timeit

df_for1 = pd.DataFrame()
for i in range(len(df)):
    if df['bool'].iloc[i] == True:
        if df['val'].iloc[i] == 5:
            df_for1 = pd.concat([df_for1, df.iloc[i]], axis=1)
            
df_for1 = df_for1.transpose()

実行結果
7.34 ms ± 2.64 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit

df_rownum = []
for i in range(len(df)):
    if df['bool'].iloc[i] == True:
        if df['val'].iloc[i] == 5:
            df_rownum.append(i)
    
df_for2 = df.iloc[df_rownum]

実行結果
1.66 ms ± 411 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

100個のデータを持つデータフレームの場合は4倍程度の速度アップになりました。

さらに10000個のデータに増やしてみたところ、concatを使う方法は1.16秒、行番号を使う方法は168ミリ秒と約7倍の速度アップになっています。

ということで良かったら試してみてください。

次回はSQLite3でデータの最後の行を取得する方法を紹介します。