python - 基于跨列的多个条件有效选择 Pandas 数据框中的行

Question

我正在尝试根据条件创建一个新的熊猫数据框。这是原始数据框：

        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

我写了这个“几乎”工作的非常复杂的启动器......但是我在处理满足 topic1 和 topic2 条件的行时遇到了问题。有什么更有效和正确的方法来做到这一点？

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

标签： pythonpandasloopsif-statementconditional-statements

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"])