如何平衡数据集

以下是该脚本的通用版本,以便您可以基于一行和要在该行中平衡的一些值来平衡任何数据集:

# 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)

现在,可能有更快的方法来做到这一点——鼓励其他人发布自己的解决方案。