How to balance a dataset

user_5049813 0
Viewed 36

I have a CSV file that has rows with a column called "worked", and I want to balance the amount of rows where "worked" is true/false. (Have them both have the same number of rows.)

I had a previous script for balancing a dataset when the column was "label" and the values were binary 0 or 1, but I'm unsure quite how to extend that to this case, or, even better, generalize it.

My old script:

# balance the dataset so there are an equal number of 0 and 1 labels

import random
import pandas as pd

INPUT_DATASET = "input_dataset.csv"
OUTPUT_DATASET = "output_dataset.csv"

LABEL_COL = "label"

# load the dataset
dataset = pd.read_csv(INPUT_DATASET)

# figure out the minimum number of 0s and 1s
num_0s = dataset[dataset[LABEL_COL] == 0].shape[0]
num_1s = dataset[dataset[LABEL_COL] == 1].shape[0]
min_num_rows = min(num_0s, num_1s)
print(f"There were {num_0s} 0s and {num_1s} 1s in the dataset - the kept amount is {min_num_rows}.")

# randomly select the minumum number of rows for both 0s and 1s
chosen_ids = []
for label in (0, 1):
    ids = dataset[dataset[LABEL_COL] == label].index
    chosen_ids.extend(random.sample(list(ids), min_num_rows))

# remove the non-chosen ids from the dataset
dataset = dataset.drop(dataset.index[list(set(range(dataset.shape[0])) - set(chosen_ids))])

# save the dataset
dataset.to_csv(OUTPUT_DATASET, index=False)
pandas dataset csv python
1 Answers

Here's a generalized version of the script so that you can balance any dataset based on a row and some values that you want to balance within that row:

# balance the given dataset based on a column and values in that column to balance

import random
import pandas as pd

RANDOM_SEED = 97

INPUT_DATASET = "input_dataset.csv"
OUTPUT_DATASET = "output_dataset.csv"

BALANCE_COL = "working"
VALUES = [True, False]

# set the random seed for reproducibility
random.seed(97)

# load the dataset
dataset = pd.read_csv(INPUT_DATASET)

# figure out the minimum number of the values
value_counts = []
for value in VALUES:
    value_counts.append(dataset[dataset[BALANCE_COL] == value].shape[0])
min_num_rows = min(value_counts)
for index, value in enumerate(VALUES):
    print(f"There were {value_counts[index]} {value}s in the dataset - the kept amount is {min_num_rows}.")

# randomly select the minumum number of rows each of the values
chosen_ids = []
for label in VALUES:
    ids = dataset[dataset[BALANCE_COL] == label].index
    chosen_ids.extend(random.sample(list(ids), min_num_rows))

# remove the non-chosen ids from the dataset
dataset = dataset.drop(dataset.index[list(set(range(dataset.shape[0])) - set(chosen_ids))])

# save the dataset
dataset.to_csv(OUTPUT_DATASET, index=False)

Now, there may be faster ways of doing this - others are encouraged to post their own solutions.

user_50498130
Related