Python: pandas で対応関係を渡して値を変換する

例えばデータセットの中のカラムが文字列型なんかで入っていると、それを数値に直したくなることはよくあると思う。今回はそれを pandas でやる方法について書く。結論から先に書くと Series オブジェクトにある map() メソッドを使うと上手くいく。

使った環境は次の通り。

$ python -V
Python 3.6.4
$ pip list --format=columns | grep -i pandas
pandas             0.22.0

まずは pandas をインストールして Python の REPL を起動しておく。

$ pip install pandas
$ python

サンプルになる DataFrame オブジェクトを用意する。各行には何かの商品のグレードと価格に関する情報が入っているイメージ。

>>> import pandas as pd
>>> data = [
...     (1, 'C', 1000),
...     (2, 'B', 2500),
...     (3, 'A', 5000),
...     (4, 'S', 10000),
... ]
>>> columns = ['id', 'grade', 'price']
>>> df = pd.DataFrame(data, columns=columns)
>>> df
   id grade  price
0   1     C   1000
1   2     B   2500
2   3     A   5000
3   4     S  10000

この中で grade カラムは文字列で情報が格納されているため扱いにくい。そこで、このカラムに入っている内容を数値に変換したいと考える。

>>> df['grade']
0    C
1    B
2    A
3    S
Name: grade, dtype: object

対応関係については次のように S ~ C に対して数値を割り振ることにする。

>>> grade_mapping = {
...     'C': 1,
...     'B': 2,
...     'A': 3,
...     'S': 4,
... }

次の通り DataFrame オブジェクトの中で各カラムは Series オブジェクトとして管理されている。

>>> type(df['grade'])
<class 'pandas.core.series.Series'>

この場合、次のように変換したいカラムに対して Series#map() メソッドを使うと上手くいく。

>>> df['grade'].map(grade_mapping)
0    1
1    2
2    3
3    4
Name: grade, dtype: int64

変換後のカラムを元の DataFrame の別のカラムと一緒に使いたいときは、次のようにすると良い。以下のコードでは既存の DataFrame に grade_norm という名前で数値に正規化したカラムを追加している。

>>> df.assign(grade_norm = df['grade'].map(grade_mapping))
   id grade  price  grade_norm
0   1     C   1000           1
1   2     B   2500           2
2   3     A   5000           3
3   4     S  10000           4

上記は DataFrame を非破壊的に操作する (新しいオブジェクトを作る) やり方になる。もし、既存の DataFrame を破壊的に更新したいときは次のように既存のカラムに代入してしまえば良い。

>>> df['grade'] = df['grade'].map(grade_mapping)
>>> df
   id  grade  price
0   1      1   1000
1   2      2   2500
2   3      3   5000
3   4      4  10000