我有一個資料集(df)如下
Company Col1 Col2 Output
AB 10 20 1
AB 20 22 1
AB 14 12 0
XZ 33 22 1
XZ 43 62 0
我想 train_test_split 資料,這樣如果一家公司在測驗集中,它根本不應該在訓練集中。我的意思是,如果第一行 ( AB, 10, 20,1) 在測驗集中,那么第二行 ( AB, 20,22,1) 也應該在測驗集中。我知道 stratify 會 stratify=df[["Name"] 會做與我想要的完全相反的事情。是否有任何內置功能可以這樣做?
PS公司列是字串
uj5u.com熱心網友回復:
這可能有點冗長而不是通用函式,但這種方法可能對您有用:
counts = df.groupby("Company").count()["Output"]
frac = 0.8 # Fraction of the training table, will only be approximated
train_companies = []
i = 0
c = 0
total_count = counts.values.sum()
train_count = total_count * frac
while(c < train_count):
train_companies.append(counts.index[i])
c = c counts.values[i]
i = i 1
c = c counts.values[i]
df_train = df[df['Company'].isin(train_companies)]
df_test = df[~df['Company'].isin(train_companies)]
轉載請註明出處,本文鏈接:https://www.uj5u.com/houduan/518338.html
上一篇:根據條件過濾具有重復值的資料框
