Pandas中按ID+col分组转置DataFrame的实现问题

Great start with adding the cumcount column—you're already halfway there! The issue with your current approach is that you're just tagging rows with their sequence number, but not collapsing the duplicate ID+col rows into one. We'll use that count value to pivot the long DataFrame into the wide format you need.

Step 1: Clarify the Original DataFrame #

First, let's fix the formatting of your input data for clarity. Your original DataFrame looks like this:

Step 2: Add the Count Column (Your Existing Code) #

You already have this part right—we tag each row in the ID+col group with its sequence number:

import pandas as pd

# Assume this is your original DataFrame
df = pd.DataFrame({
    'ID': [1, 1, 1, 1],
    'col': ['A', 'A', 'B', 'C'],
    'col2': [50, 52, 45, 18],
    'col3': ['S', 'M', 'N', 'S'],
    'col4': [1, 4, 8, 7]
})

# Add sequence count per ID+col group
df['COUNT'] = df.groupby(['ID', 'col']).cumcount() + 1

Step 3: Pivot to Wide Format #

Use pivot_table to collapse each ID+col pair into a single row, using the COUNT value to split columns like col2 into colA/colB:

# Pivot the DataFrame
pivoted = df.pivot_table(
    index=['ID', 'col'],
    columns='COUNT',
    values=['col2', 'col3', 'col4'],
    aggfunc='first'  # Use first since we only have max 2 rows per group
)

Step 4: Rename Columns to Match Your Expected Output #

The pivoted DataFrame has multi-level columns—we'll rename them to colA, colB, etc.:

# Define a function to map multi-level columns to your desired names
def rename_columns(col_tuple):
    orig_col, count = col_tuple
    if orig_col == 'col2':
        return f'col{"AB"[count-1]}'
    elif orig_col == 'col3':
        return f'col{"CD"[count-1]}'
    elif orig_col == 'col4':
        return f'col{"EF"[count-1]}'

# Apply the renaming
pivoted.columns = pivoted.columns.map(rename_columns)

Step 5: Reset Index and Replace NaN with NULL #

Finally, convert the index back to columns and replace missing values with NULL as requested:

# Reset index to make ID and col regular columns
result = pivoted.reset_index()

# Replace NaN with 'NULL' (adjust if you need actual None values instead)
result = result.fillna('NULL')

Final Output #

Running the above code will give you exactly the format you want:

The pivot operation is the key here—it takes the multiple rows per ID+col group and collapses them into one, using the count value to spread the values across new columns.

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