Randomly Split Dataframe into smaller chunks based on multiple (and similar) string conditions (rows can only used once)

user_14285047 0
Viewed 231

I have a data frame with 800 items (rows), and each row is located in a different area. The areas are: Allston, Boston, Brighton, Fenway, Brookline, Cambridge, Newton.

Example Pandas Dataframe:

       area       price      location                   Bedroom
1      boston     3074        1 Devonshire Place        1
2      boston     3310       72 Staniford Street        2
3      allston    1825  1156 Commonwealth Avenue        1
4      cambridge  3895         39 Clinton Street        3
5      fenway     2325     98 Queensberry Street        1

I try to divide the rows of this data frame into three groups RANDOMLY:

  • Group A has 60% of the rows from the data frame and could only have the following areas: 'Allston, 'Boston, 'Brighton's, 'Fenway, 'Brookline, ' Cambridge, 'newton'

  • Group B has 30% of the rows from the data frame: and could only have the following areas: 'Allston, 'Boston, ' Brighton's, 'Fenway.

  • Group C has 10% of the rows from the data frame and could only have the following areas: 'Boston, ' Brighton, ' Fenway

Every item/row can only be distributed once. It does not matter if some of the areas are not covered in one of the groups. If group C only has items that are in 'Boston, and/or Brighton, that would be okay. But group C cannot have an item that is in Newton, for instance.

I have tried dataframe.sample(), np.split(), np.random.choice(), however with all these techniques, rows get duplicated. I plan to write a loop so that the randomly selected rows will be different every time the groups get created.

Any idea on how to solve it?

conditional-statements dataframe random python
1 Answers

Here's code for this specific case. I don't think it can be generalized much.

The idea is from the areas that are allowed in group C, sample 0.1 * 800 = 80 of them and mark them as belonging to group C. Then from the unmarked rows that have areas allowed in group B, select 240 of them and mark them as belonging to group B. All the rest must be in group A.

import pandas as pd
import random

allowed = {
    'A':"Allston Boston Brighton Fenway Brookline Cambridge Newton".split(),
    'B':"Allston Boston Brighton Fenway".split(),
    'C':"Boston Brighton Fenway".split()
}

weight = {
    'A':0.6,
    'B':0.3,
    'C':0.1
}

# create random areas that meet the requirements 10% group C, 30% group B and 60% group A
rows = []
for area in random.choices(list(weight.keys()), weights=weight.values(), k=800):
    rows.append(random.choice(allowed[area]))

# create a dummy data frame
df = pd.DataFrame({'areas':rows,
                   'price':[random.randrange(1000, 5000) for _ in range(len(rows))]})

# add a column for the group, set to '' to indicate unassigned
df['group'] = ['']*len(rows)

for group in 'CBA':
    # Select rows that are not assigned to a group and that have areas that are
    # allowed for the current group. Then randomly sample the selected rows.
    xs = df[(df.group=='') & df['areas'].isin(allowed[group])].sample(n=int(len(rows)*weight[group]))

    # Mark the sampled rows with the group
    df.loc[xs.index,'group'] = group

    # this just to see what's happening
    print(group, len(xs))
    print(df.head())
    print()

The end result is the DataFrame has a column 'group' with a randomly assigned value according to the given constrains.

user_58288830
Related