Keep the Mix: stratify

When one answer is rare, a random split can share it out badly. Only about 11% of the shop's customers spent more than 400. Split the 300 customers at random and the 75 test rows may hold 7 of these big spenders one time and 13 the next. Then the test part no longer looks like the real customers, and the training part tilts the other way.

The mean of a column of 0s and 1s is the share of 1s. So y_test.mean() gives back the share of big spenders in the test part, and y_train.mean() the share in the training part. Compare them with the share in all rows to see how fair a split is.

train_test_split takes one more argument for this: stratify=y. With it, the rows where y is 1 and the rows where y is 0 are shared out separately, a quarter of each to the test part, so both parts keep the same mix, as close as whole rows allow. For example, to keep the share of customers who returned an item:

r = df["returned_item"]
X_tr, X_te, r_tr, r_te = train_test_split(X, r, test_size=0.2, stratify=r, random_state=1)
print(round(r_tr.mean(), 3), round(r_te.mean(), 3))

Below the six lines the starting code prints, your script must add:

Stratified 0 | train: 0.116 | test: 0.107
Stratified 1 | train: 0.111 | test: 0.12
Stratified 2 | train: 0.116 | test: 0.107
Stratified 3 | train: 0.116 | test: 0.107
Stratified 4 | train: 0.116 | test: 0.107

Your task: the starting code marks the big spenders (spent above 400) with 1 in y and prints their share in all rows. Then it splits five times at random and prints the share in each part. Below that loop, split again with the same five seeds, this time with stratify=y, and print the shares the same way so you can compare.

Assignment
Follow the Coding Tutorial to add a second loop of stratified splits and print the share of big spenders in each part.

Hint
The new loop is the first loop with two changes: stratify=y in the split and "Stratified" at the start of the print. Indent both lines inside the for.