Pandas按组移除DataFrame末尾满足条件的行

Hey there! Let's tackle this problem step by step. Your goal is to trim each group (grouped by id) by removing rows from the end until you hit a non-NaN value in numeric, and then remove that non-NaN row too. Here's a solution that perfectly matches your requirements:

import pandas as pd
import numpy as np

# Your original DataFrame
df = pd.DataFrame({
    "id": [1, 1, 1, 2, 2, 2, 3, 3, 3, 3],
    "date": [pd.Timestamp(2002, 2, 2), pd.Timestamp(2003, 3, 3), pd.Timestamp(2004, 4, 4),
             pd.Timestamp(2005, 5, 5), pd.Timestamp(2006, 6, 6), pd.Timestamp(2007, 7, 7),
             pd.Timestamp(2008, 8, 8), pd.Timestamp(2009, 9, 9), pd.Timestamp(2010, 10, 10), pd.Timestamp(2011, 11, 11)],
    "numeric": [0.9, 0.4, 0.2, 0.6, np.nan, 0.8, 0.7, np.nan, np.nan, 0.5],
    "nominal": [0, 1, 0, 1, 0, 0, 0, 1, 1, 1]
})

def filter_group(group):
    # Get the positions (within the group) of non-NaN values in 'numeric'
    valid_positions = group['numeric'].notna().to_numpy().nonzero()[0]
    
    if len(valid_positions) <= 1:
        # If there are 0 or only 1 non-NaN value, return empty (since we need to remove that single non-NaN and any preceding NaNs)
        return pd.DataFrame(columns=group.columns)
    
    # Keep all rows up to and including the second-to-last non-NaN position
    last_keep_position = valid_positions[-2]
    return group.iloc[:last_keep_position + 1]

# Apply the filter to each group and clean up the index
result_df = df.groupby('id', group_keys=False).apply(filter_group).reset_index(drop=True)

print(result_df)

How this works: #

1. Identify valid positions: For each group, we find the 0-based positions within the group where numeric is not NaN using notna().to_numpy().nonzero()[0].
2. Handle edge cases: If a group has 0 or only 1 non-NaN value, we return an empty DataFrame — this aligns with your logic, as we'd need to remove that single non-NaN row plus any trailing NaNs, leaving nothing.
3. Trim the group: For groups with 2+ non-NaN values, we keep all rows up to and including the second-to-last non-NaN position. This automatically takes care of:
   • Removing the last non-NaN row (as required)
   • Dropping all rows after it (which are NaNs)
   • Cutting any NaN rows between the second-to-last and last non-NaN rows

Verification against your examples: #

• id=1: Valid positions are [0,1,2]. We keep up to position 1, so rows 0 and 1 are retained (only the last row with 0.2 is removed).
• id=2: Valid positions are [0,2]. We keep up to position 0, so only the first row with 0.6 is retained (removing the NaN row and the last non-NaN row with 0.8).
• id=3: Valid positions are [0,3]. We keep up to position 0, so only the first row with 0.7 is retained (removing the two NaN rows and the last non-NaN row with 0.5).

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