如何基于伯努利变量从DataFrame中精准删除500条高Y值优先的行

一、生成恰好包含500个True的删除标记向量 #

你已经按Y列升序排序，行索引越大对应Y值越大，要实现"Y越大被删除概率越高"且恰好删除500行，加权随机抽样是最可靠的方案，比伯努利方法更直接稳定：

import numpy as np
import pandas as pd

# 假设已完成排序：df_sorted = df.sort_values("mycolumn")
total_rows = len(df_sorted)
need_delete = 500

# 生成与行索引正相关的权重（索引越大，权重越高，对应Y越大被选中概率越高）
# 用索引+1避免0权重，保证最小Y的行也有被选中的可能
weights = np.arange(1, total_rows + 1)

# 按权重随机抽取500个不重复的行位置
delete_positions = np.random.choice(total_rows, size=need_delete, replace=False, p=weights / weights.sum())

# 构造布尔标记向量
delete_index = np.zeros(total_rows, dtype=bool)
delete_index[delete_positions] = True

如果坚持想用伯努利思路，可通过截断调整保证恰好500个True，但效率较低：

p_i = np.linspace(0, 1, total_rows)
delete_index = np.random.binomial(1, p_i, size=total_rows).astype(bool)

# 循环调整标记向量
while delete_index.sum() != need_delete:
    if delete_index.sum() > need_delete:
        # 多了则去掉「Y较小、本不该被选中」的行（按索引升序取多余的True）
        extra_trues = np.sort(np.where(delete_index)[0])
        delete_index[extra_trues[:delete_index.sum() - need_delete]] = False
    else:
        # 少了则补充「Y较大、本该被选中」的行（按索引降序取缺少的False）
        missing_falses = np.sort(np.where(~delete_index)[0])[::-1]
        delete_index[missing_falses[:need_delete - delete_index.sum()]] = True

二、删除DataFrame对应行 #

有两种简单方式实现删除：

1. 布尔反选保留剩余行（最直观）：

df_remaining = df_sorted[~delete_index]

2. 使用drop方法删除指定行：

# 基于排序后的DataFrame索引删除
df_remaining = df_sorted.drop(df_sorted.index[delete_index])

若需恢复原始DataFrame的索引顺序，可在删除后执行df_remaining = df_remaining.sort_index()。

内容的提问来源于stack exchange，提问作者Programming Noob