如何用Pandas在分组中选取discharge_location为Readmit的前一行并生成字典？

Let's break down how to solve this problem: we need to group our DataFrame by ID, locate the row immediately before any row where discharge_location is "Readmit", then store those rows' ID and discharge_location in a dictionary.

Step 1: Prepare the Sample Data #

First, let's recreate your sample DataFrame and clean up the date columns (since we need proper chronological order to find the correct "previous" row):

import pandas as pd

# Build the sample DataFrame
data = {
    'ID': [20, 20, 20, 30, 30, 30, 40, 40, 40],
    'admit': ['3-4-2018', '2-2-2018', '2-5-2018', '1-2-2018', '1-15-2018', '1-20-2018', '1-20-2018', '1-20-2018', '1-20-2018'],
    'discharge': ['3-6-2018', '2-6-2018', '2-23-2018', '2-3-2018', '1-18-2018', '1-24-2018', '1-24-2018', '1-24-2018', '1-24-2018'],
    'discharge_location': ['Home', 'Home', 'Readmit', 'Rehab', 'Readmit', 'Home', 'Home', 'Home', 'Home']
}

df = pd.DataFrame(data)

# Convert date strings to datetime objects for accurate sorting
df['admit'] = pd.to_datetime(df['admit'], format='%m-%d-%Y')
df['discharge'] = pd.to_datetime(df['discharge'], format='%m-%d-%Y')

Step 2: Sort Rows Chronologically per ID #

The original row order might not reflect the actual sequence of admissions, so we sort each ID group by admission date to ensure we're grabbing the correct preceding row:

# Sort by ID, then by admission date to maintain chronological order
df_sorted = df.sort_values(['ID', 'admit']).reset_index(drop=True)

Step 3: Identify Target Preceding Rows #

We'll use groupby and shift() to mark the row immediately before each "Readmit" entry:

• First, flag all rows where discharge_location is "Readmit"
• Shift this flag forward by one row within each ID group (so the flag lands on the row before the Readmit)
• Filter the DataFrame to keep only these flagged rows

# Flag rows where discharge_location is "Readmit"
df_sorted['is_readmit'] = df_sorted['discharge_location'] == 'Readmit'

# Shift the flag to the previous row within each ID group (fill NaN with False for edge cases)
df_sorted['is_target_row'] = df_sorted.groupby('ID')['is_readmit'].shift(1).fillna(False)

# Extract only the target rows
target_rows = df_sorted[df_sorted['is_target_row']]

Step 4: Convert to the Desired Dictionary #

Finally, we convert the filtered rows into a dictionary mapping ID to discharge_location:

# Create the result dictionary
result_dict = target_rows.set_index('ID')['discharge_location'].to_dict()

print(result_dict)
# Output: {20: 'Home', 30: 'Rehab'}

Key Notes #

• Edge Case Handling: If a group's first row is a "Readmit", it won't have a preceding row, so it's automatically excluded (thanks to fillna(False)).
• Multiple Readmits: If an ID has multiple "Readmit" entries, each corresponding preceding row will be included in the dictionary (if they exist).

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