如何从联合概率DataFrame还原原始DataFrame？

我有一个DataFrame，已为其生成联合概率表，希望借此还原出与原始数据行数完全一致的原DataFrame（而非生成随机数据）。我编写了生成联合概率表的代码及还原代码，但还原代码无法输出正确结果，且希望保留原数据中的NaN行。请问是否存在可行方法实现该需求？

def compute_joint_probability(input_csv):
    """
    Compute the joint probability for each row in the input CSV file and add it as a new column in the output CSV file.

    Args:
        input_csv (str): Path to the input CSV file.
        output_csv (str): Path to the output CSV file.

    Returns:
        None
    """
    output_csv = 'JointProb.csv'
    df = pd.read_csv(input_csv)
    N = len(df)
    # Fill NaN values with a placeholder to include them in the count
    df_filled = df.fillna('NaN')
    
    # Count the occurrences of each unique row, including rows with NaN
    row_counts = df_filled.groupby(df_filled.columns.tolist()).size().reset_index(name='P')
    row_counts[row_counts.columns[-1]] = row_counts[row_counts.columns[-1]].apply(lambda x: x / N)
    # Save the counts to a new CSV file
    row_counts.to_csv(output_csv, index=False)

import random
from collections import Counter
import pandas as pd
import csv

def revert(joint_prob_csv, original_data_csv):
    # Read the joint probability CSV
    joint_prob_df = pd.read_csv(joint_prob_csv)

    # Calculate the total row count from the last column in the joint probability CSV
    N = int(joint_prob_df.iloc[:, -1].sum())

    # Remove the probability column to get the original data
    original_data_df = joint_prob_df.iloc[:, :-1]

    # Create a list to store the original rows
    original_rows = []

    # Iterate over the rows of the joint probability data
    for _, row in joint_prob_csv.iterrows():
        # Repeat each row according to its probability
        repeat_count = int(row['P'] * N)
        original_rows.extend([row[:-1]] * repeat_count)

    # Convert the list of original rows to a DataFrame
    original_data_final_df = pd.DataFrame(original_rows, columns=original_data_df.columns[:-1])

    # Save the original data with the actual row count to a new CSV file
    original_data_final_df.to_csv(original_data_csv, index=False)

    return original_data_csv


# main 
revert('JointProb.csv', 'revert.csv')  

你的还原代码存在三个核心问题：

1. 遍历对象错误：你用joint_prob_csv.iterrows()遍历，但joint_prob_csv是文件路径字符串，应该使用读取后的DataFrame对象joint_prob_df。
2. 精度丢失问题：直接用int(row['P'] * N)会因浮点数精度误差导致总行数与原始数据不一致。
3. NaN未还原：生成概率表时用'NaN'占位符，还原时未将其转回真正的pd.NA。

优化方案1：基于现有概率表修正还原代码 #

import pandas as pd

def revert(joint_prob_csv, original_data_csv):
    # 读取联合概率表
    joint_prob_df = pd.read_csv(joint_prob_csv)
    # 获取原始数据总行数（通过概率总和反推，避免浮点数误差）
    N = joint_prob_df['P'].mul(len(pd.read_csv(joint_prob_csv))).sum().round().astype(int)
    
    original_rows = []
    for _, row in joint_prob_df.iterrows():
        # 提取当前行数据（排除概率列）
        data_row = row.drop('P').to_dict()
        # 计算重复次数（概率×总行数，取整）
        repeat_count = int(row['P'] * N)
        original_rows.extend([data_row] * repeat_count)
    
    # 转为DataFrame并还原NaN
    original_data_final_df = pd.DataFrame(original_rows)
    original_data_final_df = original_data_final_df.replace('NaN', pd.NA)
    
    # 保存结果
    original_data_final_df.to_csv(original_data_csv, index=False)
    return original_data_csv

# 调用
revert('JointProb.csv', 'revert.csv')

优化方案2：修改生成代码保存计数（更准确） #

在生成联合概率表时同时保存原始出现次数，避免概率反推的精度问题：

def compute_joint_probability(input_csv):
    output_csv = 'JointProb.csv'
    df = pd.read_csv(input_csv)
    N = len(df)
    # 用占位符填充NaN
    df_filled = df.fillna('NaN')
    # 统计唯一行的出现次数并计算概率
    row_counts = df_filled.groupby(df_filled.columns.tolist()).size().reset_index(name='count')
    row_counts['P'] = row_counts['count'] / N
    # 保存包含计数和概率的表
    row_counts.to_csv(output_csv, index=False)

对应的还原代码更简洁准确：

import pandas as pd

def revert(joint_prob_csv, original_data_csv):
    joint_prob_df = pd.read_csv(joint_prob_csv)
    original_rows = []
    
    for _, row in joint_prob_df.iterrows():
        # 提取数据行（排除计数和概率列）
        data_row = row.drop(['count', 'P']).to_dict()
        # 按原始计数重复行
        original_rows.extend([data_row] * row['count'])
    
    # 转为DataFrame并还原NaN
    original_data_final_df = pd.DataFrame(original_rows)
    original_data_final_df = original_data_final_df.replace('NaN', pd.NA)
    
    original_data_final_df.to_csv(original_data_csv, index=False)
    return original_data_csv

这两种方案都能准确还原原始数据的行数和NaN值，方案2因直接使用原始计数完全避免了精度问题，是更推荐的做法。

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