基于Pandas条件匹配ID高效填充DataFrame列的技术问询

Got it, let's solve this efficiently without slow row-by-row loops—critical for large datasets! Here's a vectorized pandas approach that preserves all your existing columns while handling the mapping correctly:

Step 1: Prepare the mapping from df_input #

First, we'll reshape the POLL columns in df_input to create a clear link between the Pass_n column names and the suffixes from the POLL columns. This uses pandas' built-in reshaping functions which are optimized for speed:

import pandas as pd

# 1. Isolate POLL columns and other non-POLL columns from df_input
poll_cols = [col for col in df_input.columns if col.startswith("POLL_")]
non_poll_input_cols = [col for col in df_input.columns if not col.startswith("POLL_")]

# 2. Convert POLL columns from wide to long format
melted_input = df_input.melt(
    id_vars=non_poll_input_cols,
    value_vars=poll_cols,
    var_name="poll_column",
    value_name="pass_number"
)

# 3. Extract the suffix from POLL column names (e.g., "X" from "POLL_X")
melted_input["poll_suffix"] = melted_input["poll_column"].str.split("_").str[-1]

# 4. Format the pass column name to match df_output's "Pass_01" style (2-digit number)
melted_input["pass_column"] = "Pass_" + melted_input["pass_number"].astype(str).str.zfill(2)

# 5. Convert back to wide format to get one row per id with mapped Pass columns
poll_mapping = melted_input.pivot(
    index=non_poll_input_cols,
    columns="pass_column",
    values="poll_suffix",
    aggfunc="first"  # Use first if multiple POLL columns map to the same Pass_n
).reset_index()

Step 2: Merge the mapping into df_output #

Now we'll combine this mapping with df_output, preserving all original columns from both DataFrames:

# 1. Merge the mapping with df_output on the id column
merged_df = df_output.merge(poll_mapping, on="id", how="left")

# 2. Resolve duplicate Pass columns (original from df_output vs mapped from df_input)
# Fill empty original Pass columns with mapped values; keep existing values if present
for pass_col in [col for col in df_output.columns if col.startswith("Pass_")]:
    # _x = original df_output column, _y = mapped column from df_input
    merged_df[pass_col] = merged_df[f"{pass_col}_x"].fillna(merged_df[f"{pass_col}_y"])
    merged_df.drop([f"{pass_col}_x", f"{pass_col}_y"], axis=1, inplace=True)

# 3. Optional: Reorder columns to keep df_output's original structure first
original_output_cols = df_output.columns.tolist()
extra_input_cols = [col for col in non_poll_input_cols if col != "id"]
final_df = merged_df[original_output_cols + extra_input_cols]

Key Notes: #

• No row-wise iteration: All operations use pandas' vectorized functions, which are way faster for large datasets compared to loops.
• Preserves all columns: Both df_input's non-POLL columns and df_output's non-Pass columns are retained in the final result.
• Handles duplicates: The aggfunc="first" in pivot ensures that if multiple POLL columns map to the same Pass_n for an id, we take the first occurrence (you can adjust this to last or another aggregation if needed).
• Flexible formatting: The str.zfill(2) ensures we match the "Pass_01" style even if your pass_number values are single-digit integers.

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