如何平衡数据集

Question

我有一个CSV文件，其中有一列名为"工作"的行，我想平衡"工作"是真/假。(让它们都有相同的行数)

我以前有一个脚本，用于平衡数据集，当列为"label"值是二进制0或1，但我不确定如何将其扩展到这种情况，或者更好地推广它。

我的旧脚本:

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

Answer 1

下面是这个脚本的一个通用版本，这样你就可以基于一行和你想要在该行内平衡的一些值来平衡任何数据集:

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

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