按條件遍歷 pandas dataframe、select 行，當條件為真時，select 其他一些行，只包含唯一值

Question

我有一個很大的 (1M+) dataframe，比如

   Column A    Column B   Column C
0       'Aa'        'Ba'        14    
1       'Ab'        'Bc'        24           
2       'Ab'        'Ba'        24
...

所以基本上我有一個字符串對列表和每個字符串對的一些數字，其中該數字僅取決於 A 列。我想要做的是：

1. 遍歷 dataframe 的行
2. 對於每一行，檢查具有條件的列 C
3. 如果條件通過，select 那一行
4. 對其他 N 行進行采樣，以便每個條件通過的行總共有 N+1 行
5. 但是以某種方式對它們進行采樣，每個 N+1 組只有行，條件也通過，並且 A 列或 B 列的字符串沒有重復
6. 在不同的 N+1 組上重復並不重要，N+1 組的結果列表將比初始 df 更長的事實也不重要。 我的任務要求所有條目都在沒有重復的 N+1 組中進行處理和傳遞。

例如，條件 Column C > 15，並且 N = 5，那么對於滿足條件的行：

    Column A    Column B   Column C
78       'Ae'        'Bf'        16

例如，我們將有 N 組：

   Column A    Column B   Column C
78       'Ag'        'Br'        18
111      'Ah'        'Bg'        20
20       'An'        'Bd'        17
19       'Am'        'Bk'        18
301      'Aq'        'Bq'        32

我的初始代碼一團糟，我試過隨機抽樣行直到達到 N，檢查它們的條件，並構建一個重復的字典來檢查它們是否唯一。 然而，在數百萬長的間隔上一遍又一遍地滾動隨機數被證明太慢了。

我的第二個想法是從條件傳遞行的點向前迭代並搜索通過條件的其他行，並再次根據重復的字典檢查它們。 這開始變得更可行，但是它有問題，當到達 df 的末尾並且沒有找到 N 可行的行時，迭代必須重置到 df 的開頭。 還是覺得挺慢的。 像這樣：

    in_data = []

    for i in range(len(df)):
        
        A = df.iloc[i]['A']
        B = df.iloc[i]['B']
        
        if (condition(A)):
            
            in_data.append([A, B])
            dup_dict = {}
            dup_dict[A] = 1
            dup_dict[B] = 1
            j = i
            k = 1
            
            while (j < len(df) and k != N):

                other_A = df.iloc[j]['A']
                other_B = df.iloc[j]['B']
                
                if (condition(other_A) and
                    other_A not in dup_dict and
                    other_B not in dup_dict):
                    
                    dup_dict[other_A] = 1
                    dup_dict[other_B] = 1
                    in_data.append([other_A, other_B])
                    k += 1
                
                j += 1
                
                if (j == len(df) and k != N):
                    
                    j = 0
     
    return in_data

我最近的想法是通過 apply() 以某種方式實現它，但它開始變得太復雜了，因為我無法弄清楚如何在 apply() 中正確索引 df 並向前迭代，然后如何做重置技巧。

因此，必須有一個更精簡的解決方案。 哦，原來的 dataframe 更像是 ~60M 長，但它通過多處理在可用的 cpu 核心之間拆分和分布，因此尺寸/任務更小。

編輯：條件是動態的，即列 C 與每次檢查中的隨機數進行比較，因此不應預先屏蔽。

編輯2：一些錯別字。

Answer 1

如果我有這個權利，你是對的

data = [
["Ag", "Br", 18],
["Ah", "Bg", 20],
["An", "Bd", 17],
["Am", "Bk", 18],
["Aq", "Bq", 32],
"Aq", "Aq", 16],
]
df = pd.DataFrame(data=data, columns=['A', 'B', 'C'])

temp_df = df[(df.C > 14) & (df.A != df.B)] # e.g. condition_on_c = 14

# get the first row to sample
initial_row_index = temp_df.sample(1, random_state=42).index.values[0]
output = temp_df[temp_df.index != initial_row_index].sample(N, replace=True)
# sample = True means with replacement so you may get dup rows (definitely if N > len(temp_df) - 1
output = pd.concat([temp_df.loc[[initial_row_index]], output])

# if N = 5 we get 
    A   B   C
1  Ah  Bg  20 # initial row
3  Am  Bk  18
4  Aq  Bq  32
2  An  Bd  17
4  Aq  Bq  32
4  Aq  Bq  32

您可以看到原始索引是您正在采樣的數據框中的原始索引。 所以你可以重置這個索引。