如何用Pythonic方式基于Pandas实现按组内剩余值长度复制行？

Hey folks! Let's tackle this problem where we need to duplicate rows in a pandas DataFrame based on how many elements are left in their group, while also adding a new column z that references those remaining elements.

Example Input & Output #

First, let's recap the problem with the given data:

Input DataFrame:

import pandas as pd

df_inp = pd.DataFrame(
    {
        "x": ["a", "a", "a", "a", "b", "b", "b", "b"],
        "y": [1, 2, 3, 4, 5, 6, 7, 8],
    }
)

Expected Output:

df_out = pd.DataFrame(
    {
        "x": ["a", "a", "a", "a", "a", "a", "b", "b", "b", "b", "b", "b"],
        "y": [1, 1, 1, 2, 2, 3, 5, 5, 5, 6, 6, 7],
        "z": [2, 3, 4, 3, 4, 4, 6, 7, 8, 7, 8, 8]
    }
)

Pythonic Implementation #

Here's a clean, concise way to achieve this using pandas groupby + list comprehensions—this is super readable and efficient for most cases:

import pandas as pd

df_inp = pd.DataFrame(
    {
        "x": ["a", "a", "a", "a", "b", "b", "b", "b"],
        "y": [1, 2, 3, 4, 5, 6, 7, 8],
    }
)

# Generate the output DataFrame using a list comprehension
df_out = pd.DataFrame([
    {"x": group_key, "y": row["y"], "z": z_val}
    for group_key, group in df_inp.groupby("x")
    for idx, row in group.iterrows()
    for z_val in group.iloc[idx+1:]["y"]
])

# Optional: Reset index for clean output
df_out = df_out.reset_index(drop=True)

print(df_out)

How It Works #

Let's break down the logic step by step:

• Group by x: We split the original DataFrame into groups where each group shares the same x value (e.g., all "a"s and all "b"s).
• Iterate over each group: For every group, we loop through each row along with its index.
• Generate duplicate rows: For each row, we take all the y values that come after it in the group (using group.iloc[idx+1:]["y"]), and create a new row for each of these values—repeating the current row's y and setting the subsequent y as z.

Alternative: Using apply (For More Complex Logic) #

If you prefer a more explicit approach (great if you need to extend the logic later), you can use groupby.apply():

def process_group(group):
    rows = []
    group_length = len(group)
    # Skip the last row since it has no elements after it
    for idx in range(group_length - 1):
        current_row = group.iloc[idx]
        # Get all y values after the current row
        subsequent_ys = group.iloc[idx+1:]["y"].tolist()
        # Add a new row for each subsequent y value
        for z_val in subsequent_ys:
            rows.append({
                "x": current_row["x"],
                "y": current_row["y"],
                "z": z_val
            })
    return pd.DataFrame(rows)

df_out = df_inp.groupby("x", group_keys=False).apply(process_group).reset_index(drop=True)

This does the same thing as the list comprehension, but wraps the logic in a function for better readability when handling more complex group operations.

Verification #

Running either of these code snippets will produce exactly the df_out you're looking for—you can confirm by printing the result or using pd.testing.assert_frame_equal() to compare.

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