はじめに

この記事は、pandasを使ってデータフレームから特定の条件を満たす行を抽出する方法についてまとめたものです。
行や列の抽出に使える方法は複数あります。今回は三種類の方法を紹介し、それらの違いについても触れようと思います。
紹介する三つの方法とは

  1. 角括弧内に直接記述する
  2. loc関数を用いる
  3. query関数を用いる

です。

データの準備

まずはデータフレームを作成します。
学籍番号と徒歩通学にかかった時間、かばんの有無についてのデータだと思ってください(ここは今回あまり重要ではありません)。
データフレームの名前はdfとします。
以下のコードを実行してみましょう。

import random
import pandas as pd
random.seed(0)  # シード値を固定して同じ乱数を生成
tmp = {'id': ['M012', 'F001', 'M009', 'M005', 'F003', 'M006', 'M004', 'F002', 'F004', 'M001', 'M010', 'M002'],
       'date': ['20201010'] * 12,
      'walk_time':[random.randint(5, 20) for i in range(12)],
      'bag':[random.randint(0, 1) for i in range(12)]
}
df = pd.DataFrame(tmp)
print(df)
id date walk_time bag
0 M012 20201010 17 0
1 F001 20201010 18 0
2 M009 20201010 6 1
3 M005 20201010 13 0
4 F003 20201010 20 1
5 M006 20201010 17 0
6 M004 20201010 14 0
7 F002 20201010 20 1
8 F004 20201010 16 1
9 M001 20201010 11 0
10 M010 20201010 9 1
11 M002 20201010 14 1

方法1: 条件を直接記述する

一つ目の方法を紹介します。
早速ですが、先ほど用意したデータを使って実行していきましょう。
「'walk_time'と名付けた列の値が10より小さい」という条件をみたす行だけを取り出します。以下のように記述してみましょう。

print(df[df.walk_time < 10])
#df[df['walk_time'] < 10] も可
id date walk_time bag
2 M009 20201010 6 1
10 M010 20201010 9 1

はじめに作成したデータフレームと見比べてみてください。確かに「'walk_time'と名付けた列の値が10より小さい」という条件をみたす行だけを抽出することができました。
この方法における抽出の過程を簡単に説明します。試しに以下を実行してみましょう。

print(df.walk_time < 10)
0     False
1     False
2      True
3     False
4     False
5     False
6     False
7     False
8     False
9     False
10     True
11    False
Name: walk_time, dtype: bool

上のコードでは、「'walk_time'と名付けた列の値が10より小さい」という条件をみたす行はTrue、そうでなければFalseを返すようになっています。
df[] の角括弧内にこれを記述することで、Trueである行だけが取り出されるというわけです。その結果、条件を満たす行だけが抽出されます。
この中からさらに一部の列を抽出することもできます。

print(df[df.walk_time < 10].id)
# df[df['walk_time'] < 10].id も可
2     M009
10    M012
Name: id, dtype: object

「'walk_time'と名付けた列の値が10より小さい」という条件をみたす行の「'id'と名付けた列の値」を取得することができました。

次に、条件を満たす行だけを抽出したデータから、複数の列を抽出してみましょう。

df[df.walk_time < 10][['id', 'bag']]
id bag
2 M009 1
10 M010 1

方法2: loc関数を用いる

二つ目の方法を紹介します。
loc関数を利用して「'walk_time'と名付けた列の値が10より小さい」という条件をみたす行を抽出します。

print(df.loc[df.walk_time < 10])
#df.loc[df['walk_time'] < 10] も可

条件を満たす行だけを抽出したデータから、ある列を抽出する方法についても確認しておきます。

print(df.loc[df.walk_time < 10, 'id'])
# df.loc[df['walk_time'] < 10, 'id'] も可

先ほどと同様の結果が得られました。
次に、条件を満たす行だけを抽出したデータから、複数の列を抽出してみましょう。

print(df.loc[df['walk_time'] < 10, ['id', 'bag']])

こちらも先ほどと同様の結果が得られます。
列を指定する際、列名の書き方が方法1や次に紹介する方法3と少し異なりますので注意してください。

方法3: query関数を用いる

この記事で紹介する最後の方法です。
query関数を利用して「'walk_time'と名付けた列の値が10より小さい」という条件をみたす行を抽出します。

print(df.query('walk_time < 10'))

括弧内で、条件式をすべて文字列で書きます。
括弧内にdfが出てこない点は強みの一つです。なぜなら、dfが複数回出てくると冗長かつ修正時に面倒だからです。
SQLに近い書き方であり、SQLに慣れている人にとっては三種類の中で一番読みやすいかもしれません。

条件を満たす行だけを抽出したデータから、ある列を抽出する方法についても確認しておきます。

print(df.query('walk_time < 10').id)

先ほどと同様の結果が得られました。
次に、条件を満たす行だけを抽出したデータから、複数の列を抽出してみましょう。

df.query('walk_time < 10')[['id', 'bag']]

こちらも先ほどと同様の結果が得られます。

最後に

今回は列の抽出に関する三つの方法を紹介しました。得られる結果は同じなので、一つ一つの特徴を把握することが大切です。
方法1、方法2、方法3と下に行くほど簡潔で読みやすくなっています。論理演算子のand(&)やor(|)などを用いて複数の条件を書く際には、簡潔さの違いがより感じられるでしょう。よろしければ試してみてください。

参考文献

前処理大全〜データ分析のためのSQL/R/Python実践テクニック〜