Python Pandasの辞書式順序スライスでデータのサブセットを抽出する方法

はじめに

Pandasには、インデックス位置またはインデックスラベルを使ってデータのサブセットを選択する、2つの選択手段が用意されています。本記事では、その中でも「辞書式順序スライス（lexicographical slicing）」を使ってデータのサブセットを抽出する方法を解説します。

インターネット上にはさまざまなデータセットが公開されています。今回はKaggle（kaggle.com）で公開されている映画データセットを使用します。

準備：必要な列だけを読み込む

まず、この例に必要な列だけを指定して映画データセットを読み込みます。インデックスには映画タイトル（title）を設定します。

import pandas as pd
import numpy as np

movies = pd.read_csv(
    "https://raw.githubusercontent.com/sasankac/TestDataSet/master/movies_data.csv",
    index_col="title",
    usecols=["title", "budget", "vote_average", "vote_count"]
)
movies.sample(n=5)

インデックスは必ずソートしておく

特にインデックスが文字列で構成されている場合は、あらかじめインデックスをソートしておくことを強くお勧めします。大規模なデータセットを扱う際に、その違いを実感できるはずです。

インデックスをソートしないとどうなるのか？

「コードが永遠に実行され続ける」――なんて冗談はさておき、インデックスラベルがソートされていない場合、Pandasは条件に一致するラベルを探すために、すべてのラベルを先頭から1つずつ順番に走査しなければなりません。

索引ページのないオックスフォード英辞典を想像してみてください。目的の単語を調べるのに、ページを最初から順にめくっていくしかありませんよね。インデックスがソートされていれば、目的のラベルへ一気にジャンプできます。Pandasの内部動作もまったく同じです。

それでは、まず現在のインデックスがソート済みかどうかを確認してみましょう。

# インデックスがソートされているか確認
movies.index.is_monotonic

実行結果：

False

※ 補足：新しいバージョンのPandasでは is_monotonic は非推奨となっており、代わりに is_monotonic_increasing の使用が推奨されています。

明らかに、インデックスはソートされていません。この状態で「A」で始まる映画タイトルだけを抽出してみます。SQLで書くと、次のようなイメージです。

SELECT * FROM movies WHERE title LIKE 'A%'

これをPandasの loc を使ったスライスで表現すると、次のようになります。

movies.loc["Aa":"Bb"]

しかし、このコードはエラーになります。

ValueError: index must be monotonic increasing or decreasing
...
KeyError: 'Aa'

「index must be monotonic increasing or decreasing（インデックスは単調増加または単調減少である必要があります）」というエラーが発生しました。ソートされていないインデックスに対しては、辞書式順序スライスが使えないのです。

インデックスをソートして再挑戦

そこで、インデックスを昇順にソートしてから、同じ操作をもう一度試してみましょう。

# インデックスを昇順にソート
movies = movies.sort_index()

# もう一度確認
movies.index.is_monotonic

実行結果：

True

これでデータの準備が整い、辞書式順序スライスが使えるようになりました。それでは、「A」から「B」で始まるすべての映画タイトルを選択してみます。

movies.loc["A":"B"]

292行 × 3列 の結果が得られました。

降順ソートされたデータの場合

データが降順にソートされている場合はどうなるでしょうか。この状態で通常どおりにスライスすると、空のDataFrameが返されます。データが逆順に並んでいるため、これはごく自然な結果です。

解決策は簡単です。開始と終了のラベルを逆に指定すればOKです。

movies.loc["Bb":"Aa"]

228行 × 3列 の結果が得られました。

まとめ

• 文字列のインデックスに対しては、事前にソートしておくことで高速な検索が可能になる
• ソートされていないインデックスでスライスしようとすると ValueError が発生する
• 降順ソートされたデータの場合は、開始・終了のラベルを逆に指定する