基于条件复制Pandas DataFrame行并更新状态列

Hey there! Let's work through this problem efficiently—no clunky loops required, we can leverage Pandas' built-in functions to get this done cleanly.

Step 1: Understand the Raw Data #

First, let's recap your starting DataFrame:

import pandas as pd

df = pd.DataFrame({
    'qty': [10,7,2,1],
    'status 1': [5,2,2,0],
    'status 2': [3,2,0,1],
    'status 3': [2,3,0,0]
})

Each row represents a group of items, where qty is the total number of items, and the status X columns count how many items in the group belong to each status.

Step 2: Convert to Long Format (Key Step) #

Instead of first repeating the entire row (which leaves you stuck assigning statuses later), we'll first reshape the data to a "long" format where each status and its count gets its own row. This makes it easy to expand each status to individual item rows.

# Reshape wide status columns to long format
status_long = df.melt(
    id_vars=['qty'],  # Keep qty as an identifier column
    value_vars=['status 1', 'status 2', 'status 3'],
    var_name='status_label',
    value_name='status_count'
)

# Filter out rows where no items belong to the status (count = 0)
status_long = status_long[status_long['status_count'] > 0]

After this, status_long will have rows like:

Step 3: Expand Rows and Assign Single Status #

Now we can repeat each row based on status_count (which gives us one row per item), then extract the numeric status value:

# Repeat rows for each item in the status group
expanded = status_long.loc[status_long.index.repeat(status_long['status_count'])].reset_index(drop=True)

# Extract the numeric status from the label (e.g., "status 1" → 1)
expanded['single_status'] = expanded['status_label'].str.extract(r'(\d+)').astype(int)

# Clean up to keep only the single status column (or retain other columns as needed)
final_df = expanded.drop(['qty', 'status_label', 'status_count'], axis=1)

The result final_df will have 10+7+2+1 = 20 rows, each with a single single_status value matching the item's actual status.

Handling Real-World Data with Extra Columns #

If your actual DataFrame has additional columns (like product IDs, categories, etc.), you just need to include those columns in the id_vars parameter when using melt:

# Example DataFrame with extra columns
df_real = pd.DataFrame({
    'product_id': ['A', 'B', 'C', 'D'],
    'category': ['Electronics', 'Clothing', 'Home', 'Electronics'],
    'qty': [10,7,2,1],
    'status 1': [5,2,2,0],
    'status 2': [3,2,0,1],
    'status 3': [2,3,0,0]
})

# Reshape while preserving extra columns
status_long_real = df_real.melt(
    id_vars=['product_id', 'category', 'qty'],
    value_vars=['status 1', 'status 2', 'status 3'],
    var_name='status_label',
    value_name='status_count'
)

status_long_real = status_long_real[status_long_real['status_count'] > 0]

# Expand and clean up
expanded_real = status_long_real.loc[status_long_real.index.repeat(status_long_real['status_count'])].reset_index(drop=True)
expanded_real['single_status'] = expanded_real['status_label'].str.extract(r'(\d+)').astype(int)

final_real_df = expanded_real.drop(['qty', 'status_label', 'status_count'], axis=1)

This way, all your original metadata columns are preserved alongside the single status for each item.

Why This Approach is Better Than Loops #

Loops in Pandas are slow and hard to maintain for large datasets. Using melt and repeat are vectorized operations—they’re optimized to handle data in bulk, making your code faster and more readable.

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