pandasのデータフレームにgroup_idを付与したい

はじめに

pandasで重複削除したい場合や、集約したいときにはdrop_duplicatesやgroupbyを使えばやりたいことができます。

PandasのDataFrameやSeriesで重複要素を取り除く方法
Python
Pandas の groupby の使い方

ただ、groupbyするときと同じような条件で、各groupにgroup_idを付与したい、みたいなこともたまにありますが、やりかたがわからなかったので実装してみました。(bestプラクティスじゃないかもだけど簡単に実装できた)

group_idの付与

# pandasのimport
import pandas as pd

# データフレームの用意
df = pd.DataFrame({
    'building_name': ['Aビル', 'Aビル', 'Bビル', 'Cビル', 'Bビル', 'Bビル', 'Dビル'],
    'property_scale': ['large', 'large', 'small', 'small', 'small', 'small', 'large'],
    'city_code': [1, 1, 1, 2, 1, 1, 1]
})
df

building_name	property_scale	city_code
Aビル	large	1
Aビル	large	1
Bビル	small	1
Cビル	small	2
Bビル	small	1
Bビル	small	1
Dビル	large	1

# グループobject化
group_info = df.groupby(['property_scale', 'city_code'])

# 一応中身みておく
group_info.groups

{('large', 1): Int64Index([0, 1, 6], dtype='int64'),
('small', 1): Int64Index([2, 4, 5], dtype='int64'),
('small', 2): Int64Index([3], dtype='int64')}

# こちらも見ておく
group_info.get_group(('large', 1))

building_name	property_scale	city_code
Aビル	large	1
Aビル	large	1
Dビル	large	1

# group_idの付与
df = pd.concat([
    group_info.get_group(group_name).assign(group_id=group_id)
    for group_id, group_name
    in enumerate(group_info.groups.keys())])
df

building_name	property_scale	city_code	group_id
Aビル	large	1	0
Aビル	large	1	0
Dビル	large	1	0
Bビル	small	1	1
Bビル	small	1	1
Bビル	small	1	1
Cビル	small	2	2

いちおう関数化もしておきます

import pandas as pd
from pandas.core.frame import DataFrame

def add_group_id(df: DataFrame, by: list) -> DataFrame:
    """byの値が重複しているレコードにgroup_idを付与する.

    Args:
        df (DataFrame): 任意のデータフレーム
        by (list): グループ化するカラム名

    Returns:
        DataFrame

    """
    # すでにgroup_idカラムが入っている場合はbyにgroup_idも追加する
    if 'group_id' in df.columns:
        by += ['group_id']
    group_info = df.groupby(by=by)
    new_df = pd.concat([
        group_info.get_group(group_name).assign(group_id=group_id)
        for group_id, group_name
        in enumerate(group_info.groups.keys())])
    return new_df

追記

@r_beginners さんにコメントいただき、そもそもgroupbyにgroup_id算出の機能があるみたいです。

import pandas as pd
from pandas.core.frame import DataFrame

def add_group_id(df: DataFrame, by: list) -> DataFrame:
    """byの値が重複しているレコードにgroup_idを付与する.

    Args:
        df (DataFrame): 任意のデータフレーム
        by (list): グループ化するカラム名

    Returns:
        DataFrame

    """
    # すでにgroup_idカラムが入っている場合はbyにgroup_idも追加する
    if 'group_id' in df.columns:
        by += ['group_id']
    new_df = df.assign(group_id =df.groupby(by).ngroup())
    return new_df

@nkay さんのコメントにもあるように、pd.factorize()も使えそうですね。

pandasのメソッドもっと勉強しよう。。