pandas df 中設計的分層樣本

Question

我有一個 df，其中的列代表一個層（strat）。 我想遍歷這些層並將行拉出到新的 df, df_sample 。 如果案例很少，我想拉出一個層中的所有行。

我已經嘗試了以下方法，它可以工作。 但我想知道是否有更好的解決方案來解決這個問題。 例如，當我稍后使用真正更大的數據時， pd.concat可能會很慢。

df=pd.DataFrame({'ID': range(0,120),
             'strat': ['A', 'B', 'B', 'A', 'B', 'A', 'D', 'A', 'B', 'C', 
                       'A', 'D', 'A', 'A', 'A', 'D', 'F', 'D', 'F', 'C', 
                       'B', 'A', 'A', 'C', 'A', 'A', 'B', 'D', 'B', 'C', 
                       'C', 'A', 'C', 'A', 'C', 'A', 'D', 'C', 'C', 'A', 
                       'B', 'F', 'F', 'C', 'B', 'D', 'A', 'A', 'B', 'B', 
                       'A', 'C', 'A', 'A', 'F', 'A', 'A', 'B', 'A', 'D', 
                       'C', 'B', 'B', 'A', 'B', 'C', 'B', 'A', 'D', 'B', 
                       'B', 'A', 'A', 'C', 'D', 'F', 'F', 'A', 'B', 'C',
                      'F', 'B', 'D', 'A', 'A', 'F', 'B', 'D', 'B', 'A',
                      'F', 'D', 'A', 'A', 'C', 'B', 'B', 'C', 'C', 'B',
                      'F', 'A', 'A', 'B', 'B', 'B', 'F', 'A', 'B', 'C',
                      'A', 'A', 'A', 'B', 'B', 'A', 'A', 'A', 'B', 'B']})

df_sample=pd.DataFrame()

for i in df.strat.unique():
    temp=df[df['strat']==i]
    
    if len(temp) < 21:
        strat=temp.sample(len(temp))
        
    elif len(temp) > 20:
        strat = temp.sample(frac=0.5)
        
    df_sample=pd.concat([df_sample, strat])
    

Answer 1

您可以按“ groupby ”分組並計算每個“strat”中的條目數，然后識別少於 21 個條目的層並將它們打亂。 然后取剩余的層（條目超過 20 個的層）並對其中的 50% 進行抽樣。 最后連接兩個 DataFrame：

msk1 = df.groupby('strat')['strat'].count() < 21
less_than_21 = msk1.index[msk1]
msk2 = df['strat'].isin(less_than_21)    
out = pd.concat((df[~msk2].groupby('strat').sample(frac=0.5), df[msk2].sample(msk2.sum())))

Output：

      ID strat
110  110     A
72    72     A
46    46     A
31    31     A
92    92     A
..   ...   ...
18    18     F
9      9     C
23    23     C
42    42     F
82    82     D

[82 rows x 2 columns]

Answer 2

其他解決方案可能更快。 這是另一個，以防可讀性/可維護性更重要。

def sample_stratum(stratum):
    nrows = stratum.shape[0]
    if nrows < 21:
        output = stratum.sample(nrows)
    else:
        output = stratum.sample(frac=0.5)
    return output


# Index may be retained if needed
sampled_df = df.groupby(by=['strat']).apply(sample_stratum).reset_index(drop=True)


#    ID strat
# 0   12     A
# 1    7     A
# 2   50     A
# 3   58     A
# 4    0     A
# ..  ..   ...
# 77  41     F
# 78  42     F
# 79  16     F
# 80  76     F
# 81  90     F
# [82 rows x 2 columns]

Answer 3

為具有計數的所有組創建掩碼，然后分別處理每個組：

m = df.groupby('strat')['strat'].transform('size').lt(21)
df = pd.concat((df[~m].groupby('strat').sample(frac=0.5), 
                df[m].sample(frac=1)),
                ignore_index=True)
print (df)
    ID strat
0   71     A
1   31     A
2   72     A
3   39     A
4   83     A
..  ..   ...
77  37     C
78  85     F
79  19     C
80  34     C
81  73     C

[82 rows x 2 columns]

替代解決方案：

m = df['strat'].map(df['strat'].value_counts()).lt(21)
df = pd.concat((df[~m].groupby('strat').sample(frac=0.5), df[m].sample(frac=1)))