特定の列の要素の個数を条件にして、新たな列を追加する方法について

Question

前提・実現したいこと

Anaconda, Jupyte Notebook環境でPythonの勉強をしています。
現在は読み込んだCSVファイルを元にDataFrameの様々な集計や加工を行っていますが、特定の条件で新たな列を追加する処理に苦戦しています。

元になるDataFrame(df)

ID	分類	フラグ
1	A	前
1	A	前
1	A	前
1	A	前
1	B	後
2	A	後
2	A	前
2	A	後
2	A	後
2	B	後

Dataframeに以下のような処理を行いたいと思います。

IDが同じグループの中で、分類がAの行についてフラグを参照し、

-前が3つ以上なら、同じIDのグループ全体にX
-後が3つ以上なら、同じIDのグループ全体にy
-それ以外なら、同じIDのグループ全体にz
として新たな列(新フラグ)追加する。

完成形のイメージとしては

ID	分類	フラグ	新フラグ
1	A	前	x
1	A	前	x
1	A	前	x
1	A	前	x
1	B	後	x
2	A	後	z
2	A	前	z
2	A	後	z
2	A	後	z
2	B	後	z

発生している問題・エラーメッセージ

エラーまでたどり着けていません。

該当のソースコード

コードを書く前に、

➀
-前が3つ以上なら、同じIDのグループ全体にX
-後が3つ以上なら、同じIDのグループ全体にy
-それ以外なら、同じIDのグループ全体にz
と判定する関数を作る。

➁
「ID」「分類」でグループ化したdfに➀の関数を利用する。(グループごとに処理されることを期待)

という方針を立ててコードを書いてみたものの、よく考えたら「同じIDのグループ全体に結果を返したい」という実現したいことが全く取り込まれていない内容になってしまいました。

一応いま考えているコードを記載します。

#新フラグに何を入れるか判定する関数
def decision(row):
    
    if  row['フラグ'].count('前').sum() >= 3 :
        return 'x'
    
    elif row['フラグ'].count('後').sum() >= 3 :
        return 'y'
        
    else:
        return 'z'

#dfに新フラグを追加
df['新フラグ'] = df.groupby(['ID','分類'].agg(decision,axis=1)

試したこと

コードを書く以前にどういうロジックで処理を実現すればいいのか？の構想段階で止まってしまっています。

そもそもどのように考えればやりたいことが実現できるのかアドバイスをいただきたいです。(できればコード例もいただけると幸いです)

補足情報（FW/ツールのバージョンなど）

ここにより詳細な情報を記載してください。

Answer 1

以下のような感じでよさそうです。

import pandas as pd
from io import StringIO

s = """ID,分類,フラグ
1,A,前
1,A,前
1,A,前
1,A,前
1,B,後
2,A,後
2,A,前
2,A,後
2,A,後
2,B,後"""
df = pd.read_csv(StringIO(s))

def decision(df):
    df = df[df['分類'] == 'A'] # 分類が'A'のみで絞る
    if   (df['フラグ'] == '前').sum() >= 3 :
        return 'x'
    elif (df['フラグ'] == '後').sum() >= 3 :
        return 'y'
    else:
        return 'z'

df2 = pd.DataFrame( [(id, decision(dfg)) for id, dfg in df.groupby(['ID'])], columns=['ID','新フラグ'])
print(df2)
#   ID 新フラグ
#0   1    x
#1   2    y

ret = df.merge(df2, on='ID')
print(ret)
"""
   ID 分類 フラグ 新フラグ
0   1  A   前    x
1   1  A   前    x
2   1  A   前    x
3   1  A   前    x
4   1  B   後    x
5   2  A   後    y
6   2  A   前    y
7   2  A   後    y
8   2  A   後    y
9   2  B   後    y
"""

Comments

[upa0928]

ありがとうございます！
やりたいことが実現できました！

Answer 2

groupbyとapplyとmergeを使うので、少し複雑になります。

>>> print(df)
   ID 分類 フラグ
0   1  A   前
1   1  A   前
2   1  A   前
3   1  A   前
4   1  B   後
5   2  A   後
6   2  A   前
7   2  A   後
8   2  A   後
9   2  B   後
>>>
>>> def classify(row):
...     if row['前'] >= 3: return 'X'
...     elif row['後'] >= 3: return 'Y'
...     else: return 'z'
...
>>> df2 = df.groupby(['ID','フラグ']).count().reset_index()
>>> df3 = df2.pivot(index='ID', columns='フラグ')
>>> df3.columns = ['前', '後']
>>> df3['新フラグ'] = df3.apply(classify, axis=1)
>>> df_result = df3[['新フラグ']].merge(df, on='ID')[['ID', '分類', 'フラグ', '新フラグ']]
>>> print(df_result)
   ID 分類 フラグ 新フラグ
0   1  A   前    X
1   1  A   前    X
2   1  A   前    X
3   1  A   前    X
4   1  B   後    X
5   2  A   後    Y
6   2  A   前    Y
7   2  A   後    Y
8   2  A   後    Y
9   2  B   後    Y

Comments

[upa0928]

ご回答ありがとうございます！

いったん何も考えずに自身のDataFrameを使っていただいたコードを打ち込んでみたのですがエラーになってしまいます。

理由としてあり得そうなのは、実際のデータでは他の列も存在するということです。
どういった状況でも同様のコードで実現できますでしょうか？

また、

df2 = df.groupby(['ID','フラグ']).count().reset_index()

こちらですが、IDとフラグではなく、IDと分類によるグループ化が必要な想定でしたが間違っていますでしょうか・・・？

[upa0928]

また、私が実現したいこととしては
indexの5～9はzになる想定です。

[upa0928]

↑2つ目のコメントは間違いです。yで大丈夫でした。

[ppaul]

df2 = df.groupby(['ID','フラグ']).count().reset_index()
を
df2 = df[['ID','分類','フラグ']].groupby(['ID','フラグ']).count().reset_index()
に変更してください。
また
df_result = df3[['新フラグ']].merge(df, on='ID')[['ID', '分類', 'フラグ', '新フラグ']]
は並べ替えなので
df_result = df3[['新フラグ']].merge(df, on='ID')
としておいて、あとは必要な形に並べ替えてください。

[upa0928]

ありがとうございます！
やりたいことが実現できました！