Skip to content

Instantly share code, notes, and snippets.

@AlJohri

AlJohri/stratified_train_test_split.py

Created March 1, 2016 20:52
Show Gist options
  • Save AlJohri/bd2599098e2928b4fade to your computer and use it in GitHub Desktop.
Save AlJohri/bd2599098e2928b4fade to your computer and use it in GitHub Desktop.
Download ZIP
Raw
stratified_train_test_split.py
from sklearn.cross_validation import train_test_split
df = pd.DataFrame(np.random.randn(100, 2), columns=['x1', 'x2'])
df['y'] = np.random.choice([0, 1], size=len(df), p=[0.1, 0.9])
train, test = train_test_split(df, test_size=0.2, stratify=df.y)
train.y.value_counts()
test.y.value_counts()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment