pandasのdataframeでA列が特定の条件に当てはまる場合にB列の値が上位n個までC列に特定の値を入れたい

Question

表題の通り、pythonでpandasを使っていますが、dataframeの処理で実現したいことができない状態です。

例
A列が、A1,A2,A3...というカテゴリ変数
B1列/B2列が、全データをある基準に基づいてランク付けされた順位(1,2,...,全データ数)
C列が、空欄
とします。

やりたい処理は、A列がA1のデータのB1の上位n位までのデータのC列に1を入れ、A列がA2のデータのB2の上位n位までのデータのc列に1を入れるという処理です。

df.loc[(df[B1] <= n) & (df[A] == A1), C] = 1
df.loc[(df[B2] <= n) & (df[A] == A2), C] = 1

現在、このような処理を書いているのですが、これだとB1とB2の順位がどちらもn以下の場合重複して選ばれてしまいます。
やりたい処理としては重複してほしくなく、1行目の処理はA1のデータだけでB1のランキング上位n個、2行目はA2のデータだけでB2のランキング上位n個のものを選定(=c列に1を代入)したいです。

いい方法が思いつかず、どうしたら実現できますでしょうか？

pythonのバージョンは3.7.7です

Comments

[meg_]

具体的にdataframeを提示してください。

[pariparigorilla]

こんなイメージですね。
n=3のときの例です。

index| A | B1 | B2 | C
0 |A1 | 19 | 5 |
1 |A2 | 12 | 19 |
2 |A3 | 10 | 9 |
3 |A1 | 14 | 13 |
4 |A2 | 8 | 1 | 1(B2が1で、A2の中で1番)
5 |A1 | 11 | 16 |
6 |A2 | 18 | 10 |
7 |A2 | 9 | 4 | 1(B2が4で、A2の中で2番)
8 |A1 | 5 | 14 | 1(B1が5で、A1の中で3番)
9 |A3 | 20 | 7 |
10 |A1 | 13 | 2 |
11 |A1 | 4 | 18 | 1(B1が4で、A1の中で2番)
12 |A2 | 7 | 11 |
13 |A2 | 1 | 8 | 1(B2が8で、A2の中で3番)
14 |A3 | 15 | 3 |
15 |A3 | 6 | 17 |
16 |A2 | 17 | 15 |
17 |A1 | 2 | 12 | 1(B1が2で、A1の中で1番)
18 |A1 | 16 | 6 |
19 |A2 | 3 | 20 |

[meg_]

情報は質問に追記してください。回答が付きやすくなるかと思います。

Answer 1

df.nsmallest(n, '列名')メソッドを用いると、df['列名']の上位n位を選択することができます。
pandas.DataFrame.nsmallest — pandas documentation

なので、A列がA1やA2に一致するデータフレームに対してこのメソッドを用いれば希望のインデックス（行番号）を取得することができます。

import io
import pandas as pd

txt = """index,A,B1,B2,C
0,A1,19,5,
1,A2,12,19,
2,A3,10,9,
3,A1,14,13,
4,A2,8,1,
5,A1,11,16,
6,A2,18,10,
7,A2,9,4,
8,A1,5,14,
9,A3,20,7,
10,A1,13,2,
11,A1,4,18,
12,A2,7,11,
13,A2,1,8,
14,A3,15,3,
15,A3,6,17,
16,A2,17,15,
17,A1,2,12,
18,A1,16,6,
19,A2,3,20,"""
df = pd.read_csv(io.StringIO(txt), index_col='index')

このとき、

n = 3
df.loc[df[df['A'] == 'A1'].nsmallest(n, 'B1').index, 'C'] = 1
df.loc[df[df['A'] == 'A2'].nsmallest(n, 'B2').index, 'C'] = 1

df
#         A  B1  B2    C
# index
# 0      A1  19   5  NaN
# 1      A2  12  19  NaN
# 2      A3  10   9  NaN
# 3      A1  14  13  NaN
# 4      A2   8   1  1.0
# 5      A1  11  16  NaN
# 6      A2  18  10  NaN
# 7      A2   9   4  1.0
# 8      A1   5  14  1.0
# 9      A3  20   7  NaN
# 10     A1  13   2  NaN
# 11     A1   4  18  1.0
# 12     A2   7  11  NaN
# 13     A2   1   8  1.0
# 14     A3  15   3  NaN
# 15     A3   6  17  NaN
# 16     A2  17  15  NaN
# 17     A1   2  12  1.0
# 18     A1  16   6  NaN
# 19     A2   3  20  NaN

Comments

[pariparigorilla]

回答ありがとうございます！所望の挙動になりました。

Answer 2

こんな感じのコードで、目的の動作をすると思います。

# A列でグループ分けをしつつ各グループ内のB1/B2列の順位を求める
categorized_rank_b1 = df.groupby('A')['B1'].rank().astype(int)
categorized_rank_b2 = df.groupby('A')['B2'].rank().astype(int)
# A列がA1ならB1列順位、 A列がA2ならB2列順位、どちらかがn以下であれば、TrueをC列に入れる
df['C'] = (df['A'] == 'A1') & (categorized_rank_b1 <= n) | (df['A'] == 'A2') & (categorized_rank_b2 <= n)
# C列のTrue/Falseを、1/0に変換する
df['C'] = df['C'].astype(int)

Comments

[pariparigorilla]

回答ありがとうございます！所望の挙動になりました。

Answer 3

update()とrank()を使って、以下のように書けます。

n = 3

df['C'] = False
df['C'].update(df.loc[df['A'] == 'A1', 'B1'].rank() <= n)
df['C'].update(df.loc[df['A'] == 'A2', 'B2'].rank() <= n)
df['C'] = df['C'].astype(int)

update()はもとのindexにあわせて更新してくれるので、一部切り出して計算した結果をもとのデータフレームに代入したいときに、役に立ちます。