Random sample the model scores into 4 groups with a similar distribution in python-CodePudding

I have a dataset with model scores ranging from 0 to 1. The table looks like below:

| Score |
| ----- |
| 0.55  |
| 0.67  |
| 0.21  |
| 0.05  |
| 0.91  |
| 0.15  |
| 0.33  |
| 0.47  |

I want to randomly divide these scores into 4 groups. control, treatment 1, treatment 2, treatment 3. control group should have 20% of the observations and the rest 80% has to be divided into the other 3 equal sized groups. However, i want the distribution of scores in each group to be the same. How can i solve this using python?

PS: This is just a representation of the actual table, but it will have a lot more observations than this.

CodePudding user response：

You can use numpy.random.choice to set random groups with defined probabilities, then groupby to split the dataframe:

import numpy as np
group = np.random.choice(['control', 'treatment 1', 'treatment 2', 'treatment 3'],
                          size=len(df),
                          p=[.2, .8/3, .8/3, .8/3])

dict(list(df.groupby(pd.Series(group, index=df.index))))

possible output (each value in the dictionary is a DataFrame):

{'control':    Score
 2   0.21
 5   0.15,
 'treatment 1':    Score
 7   0.47,
 'treatment 2':    Score
 1   0.67
 3   0.05,
 'treatment 3':    Score
 0   0.55
 4   0.91
 6   0.33}

CodePudding user response：

I'm using lists just for illustration. For each number you throw a five-sided dice and if it's 1 it goes into control. If it's not 1 you throw a 3 sided dice (yeah, probably there's no such thing ;) ) and that determines the treatment group.

import random
list = [0.23, 0.034, 0.35, 0.75, 0.92, 0.25, 0.9]   
control = []
treatment1 = []
treatment2 = []
treatment3 = []
for score in list:
    dice = random.randint(1,5)
    print(dice, 'is dice')
    if dice == 1:
        control.append(score)
    else:
        seconddice = random.randint(1,3)
        print(seconddice, 'is second dice')
        if seconddice == 1:
            treatment1.append(score)
        elif seconddice == 2:
            treatment2.append(score)
        else: # seconddice == 3:
            treatment3.append(score)
    
print(control, 'is control')
print(treatment1, 'is treatment1')
print('and so on')

I made a short list to test and the results were

5 is dice
1 is second dice
1 is dice
4 is dice
3 is second dice
5 is dice
2 is second dice
1 is dice
5 is dice
1 is second dice
3 is dice
1 is second dice
[0.034, 0.92] is control
[0.23, 0.25, 0.9] is treatment1
and so on

The larger the data set, the better your distribution will be.

CodePudding user response：

Generate numbers:

import random
randomlist = []
for i in range(0,10):
    n = random.uniform(0,1)
    randomlist.append(n)

randomlist

Split into chunks: - so in this case:

categories = 4;
length = round(len(randomlist)/categories)

chunks = [randomlist[x:x length] for x in range(0, len(randomlist), length)]