如何将Pandas DataFrame拆分为n个Chunk，每个Chunk包含Col4全部唯一值

Got it, let's tackle this problem step by step. The core requirement is splitting your DataFrame into n chunks where each chunk contains at least one row for every unique value in Col4—in your example, those unique values are D1, D2, and D4. Randomly splitting by row count won't work here, so we need a targeted approach to ensure full coverage in each chunk.

Step 1: Define Your Example DataFrame #

First, let's recreate your sample data to work with:

import pandas as pd

data = [
    ["A1", "B3", "C2", "D4"],
    ["A2", "B4", "C3", "D2"],
    ["A5", "B2", "C1", "D1"],
    ["A4", "B3", "C2", "D4"],
    ["A5", "B4", "C3", "D2"],
    ["A2", "B3", "C4", "D1"],
    ["A3", "B2", "C3", "D4"],
    ["A1", "B4", "C3", "D2"],
    ["A4", "B2", "C1", "D1"],
]

df = pd.DataFrame(data, columns=["Col1", "Col2", "Col3", "Col4"])

Step 2: Create the Splitting Function #

This function will ensure every chunk gets rows for all unique values in your target column (Col4):

def split_df_by_col_unique(df, target_col, n_chunks):
    # Grab all unique values from the target column
    unique_vals = df[target_col].unique()
    
    # Group rows by each unique value in the target column
    value_groups = {val: df[df[target_col] == val].reset_index(drop=True) for val in unique_vals}
    
    # Initialize empty DataFrames for each chunk
    chunks = [pd.DataFrame(columns=df.columns) for _ in range(n_chunks)]
    
    # Distribute rows from each value group evenly across chunks
    for val in unique_vals:
        group = value_groups[val]
        # Split the group into n_chunks equal-ish parts using slicing
        row_splits = [group[i::n_chunks] for i in range(n_chunks)]
        # Append each split to its corresponding chunk
        for chunk_idx in range(n_chunks):
            chunks[chunk_idx] = pd.concat([chunks[chunk_idx], row_splits[chunk_idx]], ignore_index=True)
    
    return chunks

Step 3: Use the Function and Verify Results #

Let's test it with n=3 (matching your example):

# Split into 3 chunks
result_chunks = split_df_by_col_unique(df, "Col4", 3)

# Print each chunk to check coverage
for idx, chunk in enumerate(result_chunks, 1):
    print(f"df{idx}:")
    print(chunk)
    print(f"Unique values in Col4: {chunk['Col4'].unique()}")
    print("-" * 40)

How This Works #

• Group by Unique Values: First, we cluster all rows by their Col4 value, so we have separate lists for D1, D2, and D4 rows.
• Even Distribution: For each group, we split its rows evenly across the n chunks. Using slicing i::n_chunks ensures we spread rows out instead of dumping all rows for one value into a single chunk.
• Build Chunks: We append the split rows to each chunk, guaranteeing every chunk gets at least one row for each Col4 unique value.

The output will match your desired split (row order might vary slightly, but the coverage requirement is fully met).

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