Эффективный выбор строк в кадре данных Pandas на основе нескольких условий по столбцам

Question

Я пытаюсь создать новый pandas фрейм данных на основе условий. Это исходный кадр данных:

        topic1 topic2 
name1    1      4
name2    4      4
name3    4      3
name4    4      4
name5    2      4

Я хочу выбрать произвольные строки, чтобы topic1 == 4 появлялся 2 раза, а topic2 == 4 появлялся 3 раза в новом кадре данных. Как только это будет выполнено, я хочу остановить код.

bucket1_topic1 = 2
bucket1_topic2 = 3

Я написал этот довольно запутанный стартер, который «почти» работает ... Но у меня возникают проблемы при работе со строками, которые удовлетворяют условиям для и тема1, и тема2. Какой самый эффективный и правильный способ сделать это?

rows_list = []

counter1 = 0
counter2 = 0

for index,row in data.iterrows():
    if counter1 < bucket1_topic1:
        if row.topic1 == 4:
            counter1 +=1
            rows_list.append([row[1], row.topic1, row.topic2])

    if counter2 < bucket1_topic2:
        if row.topic2 == 4 and row.topic1 !=4:
            counter2 +=1
            if [row[1], row.topic1, row.topic2] not in rows_list:
                rows_list.append([row[1], row.topic1, row.topic2])

Желаемый результат, где topic1 == 4 появляется дважды, а topic2 == 4 появляется 3 раза:

        topic1 topic2 
name1    1      4
name2    4      4
name3    4      3
name5    2      4

Answer 1

Избегайте зацикливания и рассмотрите возможность произвольного перестановки строк с помощью DataFrame.sample (где frac=1 означает возврат 100% -ной доли кадра данных), а затем рассчитайте число бегущих групп, используя groupby().cumcount() , Наконец, фильтр с логическим поднабором:

df = (df.sample(frac=1)
        .assign(t1_grp = lambda x: x.groupby(["topic1"]).cumcount(),
                t2_grp = lambda x: x.groupby(["topic2"]).cumcount())
     )

final_df = df[(df["topic1"].isin([1,2,3])) | 
              (df["topic2"].isin([1,2,3])) |
              ((df["topic1"] == 4) & (df["t1_grp"] < 2)) |
              ((df["topic2"] == 4) & (df["t2_grp"] < 3))]

final_df = final_df.drop(columns=["t1_grp", "t2_grp"])